WARush

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

SRM614 Div1 Medium "CycleColoring"

問題

TopCoder Statistics - Problem Statement

今日、ボブは奇妙なグラフの塗り分けを数えることに挑戦している。そのグラフは小さな輪が繋がって大きな輪を作っている。

グラフは実線と点線の、2つのタイプのエッジを持つ。

小さな輪を形作るエッジは全て実線である。
大きな輪を形作るエッジは全て点線である。

より正確には、そのグラフはN個の輪を持つ。輪はC_0, C_1, ... C_N-1と順番に番号が付けられている。あなたはNこの要素を持つint vertexCountが与えられる。C_iはvertexCount[i]だけの頂点を持つ。C_iの頂点はv_i_0, v_i_1, ... v_i_vertexCount[i]-1と順番に番号が付けられている。輪であるため、最後の頂点は最初の頂点と繋がることに注意して欲しい。これらは全て実線のエッジで繋がっている。

あなたはN個の要素を持つ2つのint fromVertex, toVertexが与えられる。これらはC_iとC_i+1が次のように繋がっていることを表す。C_iの頂点v_i_fromVertex[i]とC_i+1の頂点v_i+1_toVertex[i]は、点線のエッジによって繋がっている。

ボブはK色で、このグラフを下記のルールに従い色を塗るときに、何パターン塗り分けることができるかを知りたい。

実線で繋がっている頂点は違う色で塗らなくてはならない。
点線で繋がっている頂点は同じ色で塗らなくてはならない。

色の塗り分けのパターン数を返せ。

制約

1 <= vertexCount.size() <= 50
3 <= vertexCount[i] <= 10^6
2 <= K <= 10^9



解法

それぞれの小さな輪において、fromとtoが同じ色で塗られたときのパターン数と、違う色で塗られたときのパターン数をDPでだす。

dp[v_i][same | dif] := fromから色を塗っていって、v_iの頂点をfromと同じ色で塗ったときのパターン数と、違う色で塗ったときのパターン数

小さな輪の頂点は全て実線で繋がっているため、前の頂点と同じにならないように色を塗っていく。
dp[v_i][same] = dp[v_i-1][dif]
dp[v_i][dif] = dp[v_i-1][same] * (N - 1) + dp[v_i-1][dif] * (N - 2)


次に、大きな輪において、toの頂点が0番目の輪のfromと同じ色で塗られたときのパターン数と、違う色で塗られたときのパターン数をDPする。
dp[C_i][same | dif] := C_0からC_i間で塗ったときに、C_iのtoがC_0のfromと同じ色で塗られたときのパターン数と、違う色で塗られたときのパターン数

先ほどの小さな輪の結果をP[i][same | dif]とすると、
dp[C_i][same] = dp[C_i-1][same] * P[C_i][same] + dp[C_i-1][dif] * P[C_i][dif] / (N - 1)
dp[C_i][dif] = dp[C_i-1][same] * P[C_i][dif] + dp[C_i-1][dif] * P[C_i][dif] * (N - 2) / (N - 1)


ソースコード

const long long MOD = 1000000007;

long long P[55][2];

class CycleColoring {
public:
    
    // a ^ b % MOD
    long long pow_mod(long long a, long long b) {
        long long r = 1;
        while (b > 0) {
            if (b & 1) r = r * a % MOD;
            a = a * a % MOD;
            b >>= 1;
        }
        return r;
    }

    int countColorings(vector <int> vertexCount, vector <int> fromVertex, vector <int> toVertex, int K) {

        int N = vertexCount.size();

        // 小さいサイクルのFrom・Toが同じ色・違う色(一色)がそれぞれ何パターンあるかを出す
        for (int i = 0; i < N; i++) {
            
            int n = vertexCount[i];
            int toi = i == 0 ? N - 1 : i - 1;
            int to = toVertex[toi] - fromVertex[i];
            if (to < 0) to += n;

            // 同じ色の場合を出す
            long long s = 1;
            long long d = 0;
            for (int j = 1; j < n; j++) {
                // 同 = 違
                long long ts = d;

                // 違 = 同 * (K - 1) + 違 * (K - 2)
                long long td = (s * (K - 1) + d * (K - 2)) % MOD;
                if (j == to) td = 0;

                s = ts;
                d = td;                
            }
            P[i][0] = d;

            // 違う色の場合を出す
            s = 1;
            d = 0;
            if (to != 0) {
                for (int j = 1; j < n; j++) {
                    // 同 = 違
                    long long ts = d;

                    // 違 = 同 * (K - 1) + 違 * (K - 2)
                    long long td = (s * (K - 1) + d * (K - 2)) % MOD;
                    if (j == to) ts = 0;

                    s = ts;
                    d = td;
                }
            }            
            P[i][1] = d;
        }

        // 大きいサイクルのFrom・Toが同じ・違う色がそれぞれ何パターンあるかを出す
        long long invK = pow_mod(K - 1, MOD - 2);
        long long s = K;
        long long d = 0;
        for (int i = 0; i < N; i++) {

            long long ss = P[i][0];
            long long dd = P[i][1];

            // 同 = 同 * 同 + 違 * 違 / (K - 1)
            long long ts = (s * ss + d * dd % MOD * invK) % MOD;
            // 違 = 同 * 違 + 違 * 同 + 違 * 違 * (K - 2)/(K - 1)
            long long td = (s * dd + d * ss + d * dd % MOD * (K - 2) % MOD * invK) % MOD;

            s = ts;
            d = td;
        }

        return (int)s;
    }
};