行列累乗をABC199 Fで実装する：遷移行列から期待値の計算まで

行列累乗を使える条件を先に確かめる

行列累乗を使えるのは、状態を固定長のベクトルで表せて、1回の更新を毎回同じ線形変換として書けるときです。判断するときは、次の順で確認します。

1. 次の状態を計算するために必要な情報を、長さが変わらないベクトルへまとめられるか。
2. 更新後の各成分が、現在の状態の成分を係数付きで足した形になっているか。つまり、現在の状態に行列を掛ければ次の状態になるか。
3. その係数が各回で同じか。変換が同じなら、反復回数が大きいときに累乗を二分して計算できる。

状態を行ベクトル x、1回の変換を行列 T と決めれば、1回後は xT、K 回後は xTK です。更新方法が回ごとに変わる場合や、現在の状態の線形結合では表せない場合は、そのまま同じ TK にはできません。K が小さいときは順に更新する方が簡単なこともありますが、大きいときは二分累乗で必要な行列積を約 log2K 回に抑えられます。

同じ2×2の例で右から掛ける向きを追う

頂点2個を辺1本で結び、両端の値を毎回その平均へ置き換える例を使います。初期値を行ベクトル x = [2, 4] とすると、1回の変換は T = [[1/2, 1/2], [1/2, 1/2]] です。この記事では状態ベクトルを行ベクトルにし、右から行列を掛けます。行列の行が変換前の成分、列が変換後の成分なので、出力の第 j 成分は Σi x[i] × T[i][j] です。

xT = [2×1/2 + 4×1/2, 2×1/2 + 4×1/2] = [3, 3]
(xT)T = [3, 3]T = [3, 3]

1回目に両方の値が平均の3になり、2回目も変わりません。この行列は T2 = T です。指数0では T0 = I とし、I = [[1, 0], [0, 1]] は単位行列です。したがって xT0 = xI = [2, 4] となり、操作が0回なら初期値のままです。この同じ例を後段でもコードの追加検算に使います。

ABC199 Fでは各頂点の期待値を更新する

AtCoder公式問題文（Graph Smoothing）では、単純無向グラフの辺から毎回1本を一様に選び、その辺の両端の値を平均に置き換える操作を K 回行います。求めるのは、操作後の各頂点の値そのものではなく、その期待値です。辺の選択は各回で独立です。

入力は次の形です。頂点番号 Xi, Yi は1始まりです。

N M K
A_1 A_2 ... A_N
X_1 Y_1
...
X_M Y_M

• 2 ≤ N ≤ 100
• 1 ≤ M ≤ N(N−1)/2
• 0 ≤ K ≤ 109
• 0 ≤ Ai ≤ 109
• 1 ≤ Xi, Yi ≤ N。グラフは単純無向グラフです。

出力は頂点1からNまでの期待値を1頂点につき1行で並べ、各値を 109 + 7 を法として表します。ここで期待値を行列に置き換える理由は、各頂点の次の期待値が、現在の全頂点の値の線形結合になるためです。

遷移行列の対角と隣接成分を作る

AtCoder公式解説の考え方に沿って、現在の期待値を行ベクトル x、1回後を y とし、y = xT を満たす行列 T を作ります。頂点 i の次数を di とします。

頂点 i に接続する辺が選ばれる確率は di/M です。その場合、自分の値 x[i] は平均によって半分だけ残ります。接続辺が選ばれない場合は値がそのままなので、x[i] 自身の係数は 1 − di/(2M) です。行ベクトルの規約では、これは対角成分 T[i][i] になります。

隣接頂点 j から i へ入る寄与は、辺 (i, j) が選ばれる確率 1/M と、平均における x[j] の係数 1/2 の積です。したがって T[j][i] = 1/(2M) です。辺は無向なので、行列には T[i][j] と T[j][i] の両方を加えます。こうしてできた T を各回で使うため、K 回後は A TK となります。

入力から出力まで含むPython参考コード

次は、公式問題の入力形式と公式解説の遷移式から作った参考実装です。行列積、遷移行列の構築、二分累乗、ベクトルへの適用を含めています。公式コードや提出・AC実績として示すものではありません。

import sys

MOD = 10**9 + 7

def mat_mul(a, b):
    n = len(a)
    c = [[0] * n for _ in range(n)]
    for i in range(n):
        for k in range(n):
            aik = a[i][k]
            if aik == 0:
                continue
            for j in range(n):
                c[i][j] += aik * b[k][j]
        for j in range(n):
            c[i][j] %= MOD
    return c

def vec_mul(v, matrix):
    n = len(v)
    return [sum(v[i] * matrix[i][j] for i in range(n)) % MOD for j in range(n)]

def main():
    it = iter(map(int, sys.stdin.buffer.read().split()))
    n = next(it)
    m = next(it)
    k = next(it)
    a = [next(it) for _ in range(n)]

    degree = [0] * n
    edges = []
    for _ in range(m):
        x = next(it) - 1
        y = next(it) - 1
        edges.append((x, y))
        degree[x] += 1
        degree[y] += 1

    inv_two_m = pow(2 * m, MOD - 2, MOD)
    transition = [[0] * n for _ in range(n)]
    for i in range(n):
        transition[i][i] = (1 - degree[i] * inv_two_m) % MOD
    for x, y in edges:
        transition[x][y] = (transition[x][y] + inv_two_m) % MOD
        transition[y][x] = (transition[y][x] + inv_two_m) % MOD

    result = [[int(i == j) for j in range(n)] for i in range(n)]
    base = transition
    while k > 0:
        if k & 1:
            result = mat_mul(result, base)
        base = mat_mul(base, base)
        k >>= 1

    answer = vec_mul(a, result)
    sys.stdout.write(chr(10).join(map(str, answer)))

if __name__ == '__main__':
    main()

degree で各頂点の次数を数え、inv_two_m で法 109 + 7 における 1/(2M) を作ります。対角には 1 − di/(2M)、各辺には両方向へ 1/(2M) を入れます。入力の頂点番号は読み込み時に0始まりへ直しています。

result は単位行列から始まり、base は T です。指数 k の最下位bitが1なら result に base を右から掛け、各周回で base を2乗して k を半分にします。処理後の result が TK なので、最後の vec_mul(a, result) が A TK の各成分を返します。K = 0 ではループを通らず、resultは単位行列のままです。

公式サンプル1と追加例で出力を確かめる

公式サンプル1

次の入力と期待出力は、公式問題文のサンプル1から加工せずに掲載しています。

3 2 1
3 1 5
1 2
1 3

3
500000005
500000008

辺 (1, 2) が選ばれると値は (2, 2, 5)、辺 (1, 3) が選ばれると (4, 1, 4) です。どちらも確率は1/2なので、期待値は (3, 3/2, 9/2) となり、法表現が上の出力と一致します。

追加検算例：N=2、K=2

ここからの2例は公式サンプルではなく、公式の遷移式から作った追加検算例です。2頂点・辺1本・初期値 (2, 4) では、1回目と2回目の値がともに (3, 3) です。行列の全要素が 1/2 なので T2 = T となり、上の2×2の追跡とコードが同じ結果になります。

2 1 2
2 4
1 2

3
3

追加検算例：K=0

辺は入力されますが操作回数は0です。単位行列との積で初期値が変わらないため、期待出力は (2, 4) のままです。

2 1 0
2 4
1 2

2
4

正しさと計算量

頂点 i の期待値への寄与は、自分自身からの 1 − di/(2M) と、隣接頂点 j からの各 1/(2M) です。行列 T はこの係数を行列積 y = xT の位置に置いているため、1回の期待値更新を正しく表します。各回で同じ辺選択規則を使うので、更新を K 回重ねた期待値は A TK です。コードの二分累乗は単位行列から必要な T の累乗を掛け合わせて TK を作り、最後に行ベクトル A の右から掛けています。よって vec_mul が出力する各成分は問題の期待値です。

mat_mul は N×N 行列の積を三重ループで計算するため、1回あたり O(N3) です。二分累乗では K のbit数だけ行列積を行うので、K ≥ 2 では全体が O(N3 log K) 時間です。K = 0 では累乗ループの行列積はありません。transition、累乗中の行列、積の一時行列はそれぞれ N×N で、辺リストも M ≤ N(N−1)/2 です。したがって追加メモリは O(N2) です。