Struct nekolib::graph::TreeCata

pub struct TreeCata<T> { /* private fields */ }

全方位木 DP。

木の catamorphism。 各頂点を根としたときのものをまとめて求める。

二項演算 $\circ: S\times S\to S$ と $\star: S\times T\to S$ を考える。 $(S, \circ)$ は単位元 $\mathrm{id}_{\circ}$ を持つモノイドとする。 木の各辺は一つずつ $T$ の値を持っているとし、辺 $(v, u)$ の値を $e_{v, u}$ とする。

頂点 $v$ が葉のとき、$f(v) = \mathrm{id}_{\circ}$ とする¹。 そうでないとき、$v$ に隣接する頂点を順に $\langle u_1, u_2, \dots, u_k\rangle$ とする。$v$ が根のとき、 $$f(v) = (f_v(u_1)\star e_{u_1, v})\circ\dots\circ(f_v(u_k)\star e_{u_k, v})$$ で定める。ただし、$f_v(u)$ は、「$v$ を取り除いた森において $u$ を含む木で $u$ を根としたもの」における $f(u)$ とする。

上図グラフにおいて、頂点 $1$ に隣接する頂点のうち、$0$ が最後に来るものとすれば $$\begin{aligned} f(1) &= f_0(1)\circ(f_1(0)\star e_{0, 1}) \\ &= f_0(1) \circ ((f_0(2)\star e_{2, 0})\star e_{0, 1}) \end{aligned}$$ のようになる。

このように定められる $f$ に対し、$f(0), f(1), \dots, f(n-1)$ を求める。

Idea

まず、根を $0$ として木をトポロジカルソートしておく。 これにより、ボトムアップの DP を単にループで行うことができ、$f(0)$ が求まる。次に、上で $f(1)$ を求めたときのように、トップダウンに DP をしながら（ボトムアップの DP での結果を利用して）残りの頂点について求める。

Implementation notes

トポロジカルソートの前処理パートは、木が同じであれば $(\circ, \star)$ に依らないので使い回しできる。

実装においては、$\circ: S\times S\to S$ を fold、$\mathrm{id}_{\circ}$ を empty、$\star: S\times T\to S$ を map と呼んでいる。

empty は葉での値（頂点数 1 の木での値）と fold の単位元に対応している。 葉の値を特別扱いしたいときは、セグ木に半群を乗せるときのように、 フラグを持たせれば対応できる（下記の root-leaf の距離の総和の例を参照）。 全体の頂点数が 1 だったときの処理を最後に分ける必要があるので注意。

各頂点における頂点の順序を気にした実装になっているため、「各頂点を根として DFS したときの post-order で各頂点を並べ、その列に基数 $b$・法 $m$ の rolling hash を適用したときの値を求めよ」といった問題も処理できるはず。 ハッシュ値 $h$ と、部分木サイズ $k$ に対して $(h, b^k\bmod m)$ とかを管理すればよさそう。

Complexity

$O(n)$ time.

Examples

use nekolib::graph::TreeCata;

let g = vec![
    vec![(1, ()), (2, ())],
    vec![(0, ()), (3, ()), (4, ()), (5, ())],
    vec![(0, ())],
    vec![(1, ())],
    vec![(1, ())],
    vec![(1, ())],
];

//      0 -- 2
//      |
// 4 -- 1 -- 3
//      |
//      5

let tc: TreeCata<_> = g.into();

// max distance
let empty = 0;
let map = |&x: &usize, _: &()| x + 1;
let fold = |&x: &usize, &y: &usize| x.max(y);
assert_eq!(tc.each_root(empty, map, fold), [2, 2, 3, 3, 3, 3]);

// sum of distance
let empty = (0, 0);
let map = |&(d, c): &(usize, usize), _: &()| (d + c + 1, c + 1);
let fold =
    |&x: &(usize, usize), &y: &(usize, usize)| (x.0 + y.0, x.1 + y.1);
assert_eq!(
    tc.each_root(empty, map, fold)
        .into_iter()
        .map(|x| x.0)
        .collect::<Vec<_>>(),
    [8, 6, 12, 10, 10, 10]
);

// (sum of root-leaf distance, # of leaves)
let empty = ((0, 0), false);
let map = |&(x, inner): &((usize, usize), bool), _: &()| {
    let (x1, x0) = if inner { x } else { (0, 1) };
    ((x1 + 1 * x0, x0), true)
};
let fold = |&x: &((usize, usize), bool), &y: &((usize, usize), bool)| {
    let (x1, x0) = x.0;
    let (y1, y0) = y.0;
    ((x1 + y1, x0 + y0), x.1 | y.1)
};
assert_eq!(
    tc.each_root(empty, map, fold)
        .into_iter()
        .map(|x| if x.1 { x.0 } else { (0, 1) })
        .collect::<Vec<_>>(),
    [(7, 4), (5, 4), (9, 3), (7, 3), (7, 3), (7, 3)]
);

use nekolib::graph::TreeCata;

let g = vec![
    vec![(1, 0), (2, 0)],
    vec![(0, 1), (3, 1), (4, 1), (5, 1)],
    vec![(0, 2)],
    vec![(1, 3)],
    vec![(1, 4)],
    vec![(1, 5)],
];

let tc: TreeCata<_> = g.into();

let empty = "".to_owned();
let map = |x: &String, c: &usize| {
    if x == "" { format!("{}: []", c) } else { format!("{}: [{}]", c, x) }
};
let fold = |x: &String, y: &String| {
    if x == "" && y == "" {
        "".to_owned()
    } else if x != "" && y != "" {
        format!("{}, {}", x, y)
    } else {
        format!("{}{}", x, y)
    }
};

let actual = tc
    .each_root(empty, map, fold)
    .into_iter()
    .enumerate()
    .map(|(i, x)| format!("{}: [{}]", i, x))
    .collect::<Vec<_>>();

assert_eq!(
    actual,
    [
        "0: [1: [3: [], 4: [], 5: []], 2: []]",
        "1: [0: [2: []], 3: [], 4: [], 5: []]",
        "2: [0: [1: [3: [], 4: [], 5: []]]]",
        "3: [1: [0: [2: []], 4: [], 5: []]]",
        "4: [1: [0: [2: []], 3: [], 5: []]]",
        "5: [1: [0: [2: []], 3: [], 4: []]]",
    ]
);

let empty = "".to_owned();
let map = |x: &String, c: &usize| format!("({} {} )", x, c);
let fold = |x: &String, y: &String| format!("{}{}", x, y);

assert_eq!(tc.each_root(empty, map, fold), [
    "(( 3 )( 4 )( 5 ) 1 )( 2 )",
    "(( 2 ) 0 )( 3 )( 4 )( 5 )",
    "((( 3 )( 4 )( 5 ) 1 ) 0 )",
    "((( 2 ) 0 )( 4 )( 5 ) 1 )",
    "((( 2 ) 0 )( 3 )( 5 ) 1 )",
    "((( 2 ) 0 )( 3 )( 4 ) 1 )",
]);

Applications

AtCoder における利用例。$(S, \circ, \mathrm{id}_{\circ}, T, \star)$ の定義と、$\langle f(0), f(1), \dots, f(n-1)\rangle$ を用いて答えを得る部分のみ載せている。

// typical90_am
let empty = (0, 0);
let map = |&x: &(usize, usize), _: &()| (x.0 + x.1 + 1, x.1 + 1);
let fold =
    |&x: &(usize, usize), &y: &(usize, usize)| (x.0 + y.0, x.1 + y.1);
let res: usize =
    tc.each_root(empty, map, fold).into_iter().map(|x| x.0).sum();

// abc220_f
let empty = (0, 0);
let map = |&(x1, x0): &(usize, usize), _: &()| (x1 + x0 + 1, x0 + 1);
let fold = |&(x1, x0): &(usize, usize), &(y1, y0): &(usize, usize)| {
    (x1 + y1, x0 + y0)
};
let res: Vec<_> =
    tc.each_root(empty, map, fold).into_iter().map(|(x1, _)| x1).collect();

// abc222_f
let empty = 0;
let map = |&x: &i64, &(d, w): &(i64, i64)| x.max(d) + w;
let fold = |&x: &i64, &y: &i64| x.max(y);
let res = tc.each_root(empty, map, fold);

// abc223_g
let empty = false;
let map = |&x: &bool, _: &()| !x;
let fold = |&x: &bool, &y: &bool| x | y;
let res =
    tc.each_root(empty, map, fold).into_iter().filter(|&x| !x).count();

// s8pc_4_d
let empty = (0.0, 0);
let map = |&x: &(f64, usize), _: &()| (x.0 + 1.0, 1);
let fold = |&x: &(f64, usize), &y: &(f64, usize)| {
    let v = x.0 * x.1 as f64 + y.0 * y.1 as f64;
    let d = x.1 + y.1;
    (v / 1.0_f64.max(d as f64), d)
};
let res = tc.each_root(empty, map, fold);

// dp_v
let empty = 1;
let map = |&x: &u64, _: &()| (x + 1) % m;
let fold = |&x: &u64, &y: &u64| x * y % m;
let res = tc.each_root(empty, map, fold);

// dp_p, abc036_d
let empty = (1, 1);
let map = |&x: &(u64, u64), _: &()| ((x.0 + x.1) % MOD, x.0);
let fold =
    |&x: &(u64, u64), &y: &(u64, u64)| (x.0 * y.0 % MOD, x.1 * y.1 % MOD);
let res: Vec<_> = tc
    .each_root(empty, map, fold)
    .into_iter()
    .map(|x| (x.0 + x.1) % MOD)
    .collect();
assert!(res.iter().all(|&x| x == res[0]));

// abc160_f
let mfb = ModFactorialBinom::new(n, MOD);
let f = |i| mfb.factorial(i);
let fr = |i| mfb.factorial_recip(i);

let empty = (0, 1, 1);
let map = |&x: &(usize, u64, u64), _: &()| {
    (x.0 + 1, fr(x.0 + 1), x.1 * x.2 % MOD * f(x.0) % MOD)
};
let fold = |&x: &(usize, u64, u64), &y: &(usize, u64, u64)| {
    (x.0 + y.0, (x.1 * y.1) % MOD, (x.2 * y.2) % MOD)
};
let res: Vec<_> = tc
    .each_root(empty, map, fold)
    .into_iter()
    .map(|x| map(&x, &()).2)
    .collect();
     
// tdpc_tree
let res =
    res.into_iter().fold(0_u64, |x, y| (x + y) % MOD) * mfb.recip(2) % MOD;

References

• https://qiita.com/Kiri8128/items/a011c90d25911bdb3ed3
  • トポロジカルソートで求める話が書かれている。
• https://fsharpforfunandprofit.com/posts/recursive-types-and-folds-1b/
  • catamorphism の話が載っている。

---

1. すなわち、頂点数 1 の木における $f$ の値が $\mathrm{id}_{\circ}$ となる。 ↩

Implementations

impl<T> TreeCata<T>

pub fn each_root<U: Clone>( &self, empty: U, map: impl FnMut(&U, &T) -> U, fold: impl FnMut(&U, &U) -> U ) -> Vec<U>

Trait Implementations

impl<T> From<Vec<Vec<(usize, T), Global>, Global>> for TreeCata<T>

fn from(g: Vec<Vec<(usize, T)>>) -> Self

Converts to this type from the input type.