Day 034-Q3 | 木の彩色数え上げ (Tree DP)

問題

$N$ 頂点の木を $K$ 色で塗る。以下2条件を共に満たす塗り方を $10^9 + 7$ で割った余りで求めよ。

隣接辺制約: 辺で結ばれた2頂点は異なる色
連結成分制約: 同じ色の頂点集合の誘導部分グラフは連結

つまり「各色が連結な部分木を形成する塗り方」を数える。

制約

$2 \le N \le 3000$

$2 \le K \le N$

辺は $N-1$ 本

入出力例

入力例 1 (パス)

出力: 18

パス $1{-}2{-}3{-}4$ を3色で連続区間に分割

入力例 2 (星型)

3 2
1 2
1 3

出力: 2

中心1を色A、葉2,3を色B → 葉どうし非連結 → NG。中心と片葉を同色にする2通りのみ

概念図: 木の分割 → 色割り当て

「各色 = 連結な部分木」⟺ 「木を $m$ 個の連結部分木に分割」 × 「$K$ 色から $m$ 色を選ぶ順列」

インタラクティブ: 辺の切断 → 成分数

パス $1{-}2{-}3{-}4$ の辺ボタンをクリックして切断する/しないを切り替えてみよう。

成分数: 1 | 切断辺: 0 本

ヒント (段階的開示)

ヒント1: 方向性

「各色が連結部分木」⟺ 「木を $K$ 個までの連結部分木に分割し、各成分に異なる色」。

ヒント2: アプローチ

木を $m$ 個に分割する方法数を $f(m)$ とすると、答えは $\sum_{m=1}^{K} f(m) \cdot P(K, m)$。$f(m)$ は木 DP で求める。

ヒント3: 木 DP 遷移

dp[v][m] = 頂点 $v$ の部分木を $m$ 個の連結部分木に分割する方法。子 $c$ をマージするとき:

辺 $(v,c)$ を切る → 成分数 $s_v + s_c$
辺 $(v,c)$ を繋ぐ → 成分数 $s_v + s_c - 1$

模範解答 (Python)

import sys
from sys import stdin

def main():
    MOD = 10**9 + 7
    data = stdin.read().split()
    idx = 0
    N, K = int(data[idx]), int(data[idx+1]); idx += 2

    adj = [[] for _ in range(N+1)]
    for _ in range(N-1):
        u, v = int(data[idx]), int(data[idx+1]); idx += 2
        adj[u].append(v); adj[v].append(u)

    # BFS で親子関係を確定
    root = 1
    parent = [0] * (N+1)
    order = []
    visited = [False] * (N+1)
    queue = [root]; visited[root] = True; head = 0
    while head < len(queue):
        v = queue[head]; head += 1
        order.append(v)
        for u in adj[v]:
            if not visited[u]:
                visited[u] = True
                parent[u] = v
                queue.append(u)

    children = [[] for _ in range(N+1)]
    for v in order[1:]:
        children[parent[v]].append(v)

    size = [1] * (N+1)
    for v in reversed(order):
        for c in children[v]:
            size[v] += size[c]

    # dp[v][m] = v の部分木を m 個の連結部分木に分割する方法
    dp = [[0] * (N+1) for _ in range(N+1)]
    for v in range(1, N+1):
        dp[v][1] = 1

    for v in reversed(order):
        for c in children[v]:
            new_dp = [0] * (N+1)
            sv_max = size[v] - size[c]
            for sv in range(1, sv_max + 1):
                if dp[v][sv] == 0: continue
                for sc in range(1, size[c] + 1):
                    if dp[c][sc] == 0: continue
                    val = dp[v][sv] * dp[c][sc] % MOD
                    new_dp[sv + sc]     = (new_dp[sv + sc] + val) % MOD       # 切る
                    new_dp[sv + sc - 1] = (new_dp[sv + sc - 1] + val) % MOD   # 繋ぐ
            for m in range(N+1):
                dp[v][m] = new_dp[m]

    f = dp[root]

    # P(K, m) = K * (K-1) * ... * (K-m+1)
    ans = 0
    perm = 1
    for m in range(1, min(K, N) + 1):
        if m == 1:
            perm = K % MOD
        else:
            perm = perm * (K - m + 1) % MOD
        ans = (ans + f[m] * perm) % MOD

    print(ans)

main()

Step-by-Step 解説

1問題変換
「各色 = 連結部分木」を「分割」 + 「色順列」の積に分解。$P(K,m) = K!/(K-m)!$

2木 DP の遷移
子 $c$ をマージするとき:

切る: 成分数 $s_v + s_c$
繋ぐ: 合併で成分数 $s_v + s_c - 1$

3葉→根の処理順
BFS 順を逆転 (reversed(order)) で葉→根を保証。

4答え集計
$\text{ans} = \sum_{m=1}^{\min(K,N)} f(m) \cdot K (K-1) \cdots (K-m+1)$

計算量

木 DP: $O(N^2)$ (全辺マージで部分木サイズの積)
答え集計: $O(K)$
合計: $O(N^2)$ ← $N=3000$ で約 $9 \times 10^6$

よくあるミス

ミス	原因	正しい書き方
繋ぐ場合の -1 忘れ	合併で成分数が1減ることを失念	`new_dp[sv + sc - 1]`
`f[0]` を使う	0 個分割が無意味	$m$ は 1 から開始
$P(K,m)$ の順序	$m > K$ で負になる	`min(K, N)` の範囲で打ち切り
DP 更新順序	親が子より先に更新	`reversed(order)` で葉→根

Day 034-Q3 — 木の彩色数え上げ (Broken Profile DP + Tree DP)

問題

制約