TopCoder SRM 614: CycleColoring

第一段階で、各々の長さの彩色パターンを DP で求めて、
第二段階で、第一段階で求めた彩色パターンを使って、同じような DP を構築する。

K 色全てを状態にするのではなく、開始地点と同じ色かそうでないかの２通りだけの状態を考えれば良い。

具体的な処理はソースコード参照。

mod の処理がややこしいので、ライブラリ化した ModInt クラスを使って、楽している。

public class CycleColoring {
    const int SameColor = 0;
    const int DifferentColor = 1;

    public int countColorings(int[] vertexCount, int[] fromVertex, int[] toVertex, int K) {
        var N = vertexCount.Length;
        // dp[n, SameColor] = n 個並んだ頂点をグラフについてを隣り合った頂点に異なる色を塗ったときに 
        //                    1番目とn番目が同じ色になるような彩色のパターン数。1番目の色は固定
        var dp = new ModInt[1000005, 2];
        dp[1, SameColor] = 1;
        dp[1, DifferentColor] = 0;
        for (int i = 2; i < dp.GetLength(0); i++) {
            var all = (dp[i - 1, DifferentColor] + dp[i - 1, SameColor]) * (K - 1); // 各々の状態から遷移先は K - 1 通り
            dp[i, SameColor] = dp[i - 1, DifferentColor];
            dp[i, DifferentColor] = all - dp[i, SameColor];
        }

        // 開始地点は、K 通りに彩色できる
        ModInt SameColorWithStart = K;
        ModInt DiffColorWithStart = 0;
        for (int i = 0; i < vertexCount.Length; i++) {
            var cycleLength = vertexCount[i];
            int enterIndex = toVertex[(i - 1 + N) % N];
            int exitIndex = fromVertex[i];

            int length1 = Math.Abs(exitIndex - enterIndex);
            int length2 = cycleLength - length1;

            // サイクルの彩色パターン数（一箇所の色は固定）
            var allColoringCyclePattern = dp[cycleLength, DifferentColor];
            // 入り口と出口が同じ色の場合の彩色パターン数
            var sameEnterAndExitColoringPattern = dp[length1 + 1, SameColor] * dp[length2 + 1, SameColor];
            // 入り口と出口が異なる色の場合の彩色パターン数
            var diffEnterAndExitColoringPattern = allColoringCyclePattern - sameEnterAndExitColoringPattern;

            ModInt nextSameColorWithStart =
                 SameColorWithStart * sameEnterAndExitColoringPattern +
                 DiffColorWithStart * diffEnterAndExitColoringPattern / (K - 1);
            ModInt nextDiffColorWithStart =
                (SameColorWithStart + DiffColorWithStart) * allColoringCyclePattern - nextSameColorWithStart;

            SameColorWithStart = nextSameColorWithStart;
            DiffColorWithStart = nextDiffColorWithStart;
        }

        return SameColorWithStart.ToInt();
    }
}

public struct ModInt {
    static private int mod = 1000000007;
    private long value;

    public ModInt(long x) {
        this.value = x;
        Normalize();
    }

    static private long RegularMod(long x, int mod) {
        if (x >= mod) return x % mod;
        if (x >= 0) return x;
        return (mod - RegularMod(-x, mod)) % mod;
    }

    static public void SetModValue(int m) {
        ModInt.mod = m;
    }

    private void Normalize() {
        this.value = RegularMod(value, mod);
    }

    public override string ToString() {
        return value.ToString();
    }

    public int ToInt() {
        return (int)this.value;
    }

    public static bool operator ==(ModInt c1, ModInt c2) {
        return c1.value == c2.value;
    }

    public static bool operator !=(ModInt c1, ModInt c2) {
        return !(c1 == c2);
    }

    public static ModInt operator +(ModInt x, ModInt y) {
        return new ModInt(x.value + y.value);
    }
    public static ModInt operator -(ModInt x, ModInt y) {
        return new ModInt(x.value - y.value);
    }
    public static ModInt operator *(ModInt x, ModInt y) {
        return new ModInt(x.value * y.value);
    }
    public static ModInt operator /(ModInt x, ModInt y) {
        return new ModInt(x.value * Inverse(y.value, mod));
    }

    static private long ExtendedGcd(long a, long b, ref long x, ref long y) {
        if (b == 0) {
            x = 1; y = 0;
            return a;
        } else {
            long d = ExtendedGcd(b, a % b, ref y, ref x);
            y -= a / b * x;
            return d;
        }
    }

    static private long Inverse(long a, long mod) {
        long x = 0, y = 0;
        if (ExtendedGcd(a, mod, ref x, ref y) == 1)
            return (x + mod) % mod;
        else
            throw new Exception("Invalid inverse " + a + " " + mod);
    }

    public static implicit operator ModInt(long x) {
        return new ModInt(x);
    }

    public override bool Equals(object obj) {
        if (obj == null) {
            return false;
        }
        return this.value.Equals(((ModInt)obj).value);
    }

    public override int GetHashCode() {
        return this.value.GetHashCode();
    }
}