Day 064-Q3 | 多項式行列累乗 + グラフカウント（FPS × 行列累乗色制約パス）

問題

$N$ 頂点の有向グラフが与えられ、辺 $(i, j)$ には「色ラベル $c_{ij} \in \{1, \ldots, K\}$」が付いている。

ちょうど $L$ ステップの有向路のうち、各色の出現回数が指定の $f_1, \ldots, f_K$（$\sum f_k = L$）に一致するものの数を $\mod (10^9+7)$ で求めよ。

制約

パラメータ	範囲
$N$	$1 \le N \le 20$
$M$	$1 \le M \le N^2$
$K$	$1 \le K \le 5$
$L$	$1 \le L \le 10^6$
$s, t$	$1 \le s, t \le N$

入出力例

入力例 1

3 4 2 3 1 3
1 2 1
2 3 1
1 3 2
2 3 2
1 2

出力例 1

$f_1=1, f_2=2$。頂点1→2→3（色1,1,2は違う）→ 1→2→3: 色(1,1) … L=3なので3辺。1→2→3 + 辺(2→3,2): 色(1,1,2)✓ / 1→3→? + 1→3: 色(2,…). 実際はサンプル答えは2。

概念図: 多変数生成関数行列

ヒント（段階的開示）

ヒント1: 方向性

行列累乗の自然な拡張。各辺に色変数 $x_{c_{ij}}$ を対応させ、行列エントリを「多変数多項式（生成関数）」にする。$A^L[s][t]$ の特定係数が答え。

ヒント2: アプローチ

行列エントリ = dict{color_tuple → count}
color_tuple = (f1, f2, ..., fK)（各色の累積出現回数）
行列積: 2つの辞書の「タプル加算 + カウント乗算」で計算
次数カット: ステップ数の合計が $L$ を超えたら無視

ヒント3: コード骨格

def mat_mul(X, Y, N, K, L):
    C = [[defaultdict(int) for _ in range(N)] for _ in range(N)]
    for i in range(N):
      for j in range(N):
        for k in range(N):
          for ka, va in X[i][k].items():
            for kb, vb in Y[k][j].items():
              if sum(ka) + sum(kb) > L: continue
              merged = tuple(ka[x]+kb[x] for x in range(K))
              C[i][j][merged] = (C[i][j][merged]+va*vb) % MOD
    return C

模範解答 (Python)

import sys
from collections import defaultdict
input = sys.stdin.readline
MOD = 10**9 + 7

def solve():
    line1 = input().split()
    N, M, K = int(line1[0]), int(line1[1]), int(line1[2])
    L, s, t = int(line1[3]), int(line1[4])-1, int(line1[5])-1

    def zero_mat():
        return [[defaultdict(int) for _ in range(N)] for _ in range(N)]

    def identity_mat():
        mat = zero_mat()
        for i in range(N):
            mat[i][i][tuple([0]*K)] = 1
        return mat

    A = zero_mat()
    for _ in range(M):
        u, v, c = map(int, input().split())
        u -= 1; v -= 1; c -= 1
        key = tuple(1 if i == c else 0 for i in range(K))
        A[u][v][key] = (A[u][v][key] + 1) % MOD

    f_list = list(map(int, input().split()))
    target = tuple(f_list)

    def mat_mul(X, Y):
        C = zero_mat()
        for i in range(N):
            for j in range(N):
                for k in range(N):
                    for ka, va in X[i][k].items():
                        sa = sum(ka)
                        for kb, vb in Y[k][j].items():
                            if sa + sum(kb) > L: continue
                            merged = tuple(ka[x] + kb[x] for x in range(K))
                            C[i][j][merged] = (C[i][j][merged] + va * vb) % MOD
        return C

    def mat_pow(M_mat, n):
        result = identity_mat()
        base = M_mat
        while n:
            if n & 1: result = mat_mul(result, base)
            base = mat_mul(base, base)
            n >>= 1
        return result

    R = mat_pow(A, L)
    print(R[s][t].get(target, 0))

solve()

Step-by-Step 解説

Step 1: 通常の行列累乗からの拡張

通常の行列累乗では $A^L[i][j]$ = 「$L$ ステップで $i$ から $j$ へのパス数」。ここでは各辺に色変数を付与し、エントリを多項式に拡張する。

Step 2: 色の多重集合の表現

$K$ 個の色の出現回数を長さ $K$ のタプルで表現。辞書 {tuple → count} が各エントリの多項式。

Step 3: 次数カット（重要な最適化）

ステップ数の合計（タプルの総和）が $L$ を超えるエントリはスキップ。これにより多項式エントリ数が爆発しない。

Step 4: 繰り返し2乗法

$A^L$ を繰り返し2乗法で計算（$O(\log L)$ 回の行列積）。$A^L[s][t][\text{target}]$ が答え。

Step 5: 計算量

処理	計算量
行列乗算 1 回	$O(N^3 \cdot P^2)$（$P$ = エントリ数）
繰り返し2乗法	$O(N^3 P^2 \log L)$
実用上	次数カットで $P$ が小さくなり現実的

よくあるミス

ミス	原因	正しい書き方
色の 0/1 インデックス混在	入力は 1-indexed	`c -= 1` してからタプル生成
次数カットなし	多項式エントリが指数的に増加	`if sa + sum(kb) > L: continue`
恒等行列の多項式エントリ	対角要素を空 dict にする	`mat[i][i][(0,...,0)] = 1`

次のステップ

発展問題: 辺に「重み関数 $w_{ij}(x) = a_{ij} x + b_{ij}$」が付いた行列累乗（多項式行列の $L$ 乗の特定係数）を NTT で $O(N^3 L \log L \log N)$ で解け。

自己評価

完全理解おおむね理解要復習