WARush

SRMの結果とか、解けた問題のコードを書いていきます

Codeforces #177 Div2 D "Polo the Penguin and Houses"

問題

http://codeforces.com/contest/289/problem/D

ある村にはN個の家があり、1~Nと番号付けされている。
家iはグループ1~Nのいずれかに所属していた。
家iのグループ番号をP(i)とする。

この村の家々を次のように訪れる事にした。今訪れている家の番号がxとすると、
次に番号P(x)の家に行く。次に番号P(P(x))の家に行く...を繰り返す。

そのように移動した時、以下の事が分かった。

1. 番号1からKのいずれかの家から出発したところ、番号1の家に訪れる事ができた。
2. 番号K+1からNのいずれかの家から出発したところ、番号1の家に訪れる事はできなかった。
3. 番号1の家から出発したところ、番号0ではない家から家へと移動し、
   番号1の家に戻る事ができた。(?)

上記の3つの状態となるような、
家のグループ番号のつけ方が何通りあるかを返せ。

制約

1 <= N <= 1000
1 <= K <= min( 8, N )


考えた事

ちょっと条件3がよく分からないが、
1の家から他の家を介さず1の家へ行く、
つまり家1のグループ番号が1でもいいみたい。

まず家K+1~Nのグループ番号は1~K以外のものを自由につけて良いので
(N-K)^(N-K)通りできる。

家1のグループ番号は1~Kを自由につけてよいのでK通り。

家2~Kは、かならず家1へ行けるように
グループ番号を付けなければいけないからやっかいだな~。

あれ、Kは8以下か...
つまり最大で家2~8の7軒をどうにかしなくちゃか。

全探索すると7^7で80万の状態があって、
それが条件を満たすかどうかの判定が7回だから
計算量は80万 * 7で560万回
いける!

事後

頂点数Kの全域木の種類数を調べればいいのか~ なるほど
Round #177 - どせいけいさんき。



ソースコード

const int MOD = 1000000007;

int N, K;
int to[10]; // to[i] 家iから次に向かうべき家(グループ番号)

int ok[10]; // -1 - 未チェック 0 - チェック中 1 - 家1に行けた

bool chk_dfs( int v ){
    ok[v] = 0;
    bool res = false;

    if( to[v] == 1 ){
        ok[v] = 1;
        return true;
    }

    if( ok[to[v]] == -1 ){
        res = chk_dfs( to[v] );
        ok[v] = (res ? 1 : 0);
    }else if( ok[to[v]] == 0 ){
        res = false;
        ok[v] = 0;
    }else{
        res = true;
        ok[v] = 1;
    }
    return res;
}

// to[]とグループ番号を付けた時に家1に必ず向かう事ができるか?
bool chk(){
    memset( ok, -1, sizeof(ok) );
    
    for( int i = 2; i <= K; i++ ){
        if( ok[i] != -1 ) continue;
        if( !chk_dfs( i ) ) return false;
    }

    return true;
}  

// 家v~Kでの、グループ番号の付け方の数を返す
int dfs( int v ){
    if( v > K ){
        return chk() ? 1 : 0;
    }

    int res = 0;
    for( int i = 1; i <= K; i++ ){
        if( i == v ) continue; // 自分の家に戻るようなのは論外
        to[v] = i;
        res += dfs( v + 1 );
        res %= MOD;
    }
    return res;
}

int main() {
    
    cin >> N >> K;

    // 家1
    long long ans = K;
    // 家2 - K
    if( K > 1 ){
        ans *= dfs( 2 );
        ans %= MOD;
    }
    // 家K+1 - N
    if( N != K ){
        for( int i = K + 1; i <= N; i++ ){
            ans *= (N - K);
            ans %= MOD;
        }
    }

    cout << ans << endl;
}