Problem:
G: Tree

問題

N 頂点からなる木が与えられます。木の各頂点には 1 から N の番号が割り振られています。また、N−1 本の辺のうち、 i\ (= 1, 2, ..., N-1) 本目の辺は頂点 u_i と頂点 v_i とを結んでいます。

この木の K 個の空でない部分グラフの組であって、各部分グラフがいずれも連結で、どの二つの相異なる部分グラフも頂点を共有しないものの個数を求めるプログラムを作成してください。ただし答えがとても大きくなることがあるため、998244353 で割ったあまりを答えてください。

なお、K 個の部分グラフからなる集合が同一のものは、K 個の部分グラフの順序が異なるものも同一視するものとします。

入力形式

N K
u_1 v_1
:
u_{N-1} v_{N-1}

制約

2 \leq N \leq 10^{5}

1 \leq K \leq min ( N, 300 )

1 \leq u_i, v_i \leq N

u_i \neq v_i

i \neq j に対して (u_i, v_i) \neq (u_j, v_j)

入力はすべて整数で与えられる。

与えられるグラフは木であることが保証される。

出力形式

答えを表す整数を一行に出力してください。998244353 で割ったあまりを出力することに注意してください。

入力例 1

3 2
1 2
1 3

出力例 1

5

以下の 5 通りがあります:

\{ 1 \} と \{ 2 \}

\{ 1 \} と \{ 3 \}

\{ 2 \} と \{ 3 \}

\{ 1, 2 \} と \{ 3 \}

\{ 1, 3 \} と \{ 2 \}

入力例 2

4 4
1 2
1 3
1 4

出力例 2

1

( \{ 1 \}, \{ 2 \}, \{ 3 \}, \{ 4 \} ) の 1 通りしかありません。

入力例 3

7 4
1 7
2 1
7 4
3 4
5 7
6 3

出力例 3

166

Solution in Rust (judge verdict: Accepted)

struct TreeDP<Edge> {
    size: usize,
    edge: Vec<(usize, usize, Edge)>,
}

impl<Edge>  TreeDP<Edge> {
    fn new(size: usize) -> Self {
        TreeDP {
            size: size,
            edge: vec![],
        }
    }
    fn add_edge(&mut self, a: usize, b: usize, c: Edge) {
        assert!(a < self.size && b < self.size && a != b);
        self.edge.push((a, b, c));
    }
    fn solve<Init, Merge, Value>(&self, root: usize, init: Init, merge: Merge) -> Value
    where Init: Fn(usize) -> Value, 
          Merge: Fn(Value, Value, &Edge) -> Value,
    {
        let size = self.size;
        let mut graph = vec![vec![]; size];
        for (i, e) in self.edge.iter().enumerate() {
            graph[e.0].push((e.1, i));
            graph[e.1].push((e.0, i));
        }
        let mut topo = vec![];
        let mut stack = vec![(root, root)];
        while let Some((v, p)) = stack.pop() {
            topo.push(v);
            if let Some(k) = graph[v].iter().position(|e| e.0 == p) {
                graph[v].remove(k);
            }
            for e in graph[v].iter() {
                stack.push((e.0, v));
            }
        }
        assert!(topo.len() == size);
        let mut dp: Vec<Option<Value>> = (0..size).map(|_| None).collect();
        for &v in topo.iter().rev() {
            let mut now = init(v);
            for &(u, k) in graph[v].iter() {
                let b = dp[u].take().unwrap();
                now = merge(now, b, &self.edge[k].2);
            }
            dp[v] = Some(now);
        }
        dp[root].take().unwrap()
    }
}

use std::io::Read;

fn run() {
    let mut s = String::new();
    std::io::stdin().read_to_string(&mut s).unwrap();
    let mut it = s.trim().split_whitespace();
    let n: usize = it.next().unwrap().parse().unwrap();
    let m: usize = it.next().unwrap().parse().unwrap();
    const MOD: u64 = 998_244_353;
    let mut solver = TreeDP::new(n);
    for _ in 1..n {
        let x: usize = it.next().unwrap().parse::<usize>().unwrap() - 1;
        let y: usize = it.next().unwrap().parse::<usize>().unwrap() - 1;
        solver.add_edge(x, y, ());
    }
    type Value = Vec<(u64, u64)>;
    type Edge = ();
    let init = |_v: usize| -> Value {
        vec![(1, 0), (0, 1)]
    };
    let merge = |a: Value, b: Value, _e: &Edge| -> Value {
        let mut next = vec![(0, 0); a.len() + b.len() - 1];
        for (i, a) in a.iter().enumerate() {
            for (j, b) in b.iter().enumerate() {
                next[i + j].0 = (next[i + j].0 + a.0 * (b.0 + b.1)) % MOD;
                next[i + j].1 = (next[i + j].1 + a.1 * (b.0 + b.1)) % MOD;
                if i > 0 {
                    next[i + j - 1].1 = (next[i + j - 1].1 + a.1 * b.1) % MOD;
                }
            }
        }
        next.truncate(m + 1);
        next
    };
    let res = solver.solve(0, init, merge);
    let ans = *res.last().unwrap();
    println!("{}", (ans.0 + ans.1) % MOD);
}

fn main() {
    run();
}