library_checker/tree/
point_set_tree_path_composite_sum_fixed_root.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::{Associative, Magma, Unital},
4    graph::TreeGraphScanner,
5    num::{One, Zero, mint_basic::MInt998244353 as M},
6    tree::MonoidCluster,
7};
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10struct Point {
11    sum: M,
12    cnt: u32,
13}
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16struct Path {
17    a: M,
18    b: M,
19    sum: M,
20    cnt: u32,
21}
22
23struct PointMonoid;
24impl Magma for PointMonoid {
25    type T = Point;
26    fn operate(x: &Self::T, y: &Self::T) -> Self::T {
27        Point {
28            sum: x.sum + y.sum,
29            cnt: x.cnt + y.cnt,
30        }
31    }
32}
33impl Unital for PointMonoid {
34    fn unit() -> Self::T {
35        Point {
36            sum: M::zero(),
37            cnt: 0,
38        }
39    }
40}
41impl Associative for PointMonoid {}
42
43struct PathMonoid;
44impl Magma for PathMonoid {
45    type T = Path;
46    fn operate(x: &Self::T, y: &Self::T) -> Self::T {
47        Path {
48            a: x.a * y.a,
49            b: x.b + x.a * y.b,
50            sum: x.sum + x.a * y.sum + x.b * M::new_unchecked(y.cnt),
51            cnt: x.cnt + y.cnt,
52        }
53    }
54}
55impl Unital for PathMonoid {
56    fn unit() -> Self::T {
57        Path {
58            a: M::one(),
59            b: M::zero(),
60            sum: M::zero(),
61            cnt: 0,
62        }
63    }
64}
65impl Associative for PathMonoid {}
66
67struct Dp;
68
69impl MonoidCluster for Dp {
70    type Vertex = M;
71    type Edge = (M, M);
72    type PointMonoid = PointMonoid;
73    type PathMonoid = PathMonoid;
74
75    fn add_vertex(point: &Point, vertex: &M, parent_edge: Option<&(M, M)>) -> Path {
76        let cnt = point.cnt + 1;
77        let subtotal = point.sum + *vertex;
78        let (a, b) = parent_edge.copied().unwrap_or((M::one(), M::zero()));
79        Path {
80            a,
81            b,
82            sum: a * subtotal + b * M::new_unchecked(cnt),
83            cnt,
84        }
85    }
86
87    fn add_edge(path: &Path) -> Point {
88        Point {
89            sum: path.sum,
90            cnt: path.cnt,
91        }
92    }
93}
94
95competitive::define_enum_scan! {
96    enum Query: usize {
97        0 => SetVertex { v: usize, x: M }
98        1 => SetEdge { e: usize, a: M, b: M }
99    }
100}
101
102#[verify::library_checker("point_set_tree_path_composite_sum_fixed_root")]
103pub fn point_set_tree_path_composite_sum_fixed_root(reader: impl Read, writer: impl Write) {
104    prepare_io!(reader, writer);
105    sc!(n,
106        q,
107        value: [M; n],
108        (graph, edges): @TreeGraphScanner::<usize, (M, M)>::new(n));
109
110    let top_tree = graph.static_top_tree(0);
111    let mut dp = top_tree.dp::<Dp>(value, edges);
112
113    for _ in 0..q {
114        sc!(query: Query);
115        match query {
116            Query::SetVertex { v, x } => {
117                dp.set_vertex(v, x);
118                pp!(dp.fold_all().sum);
119            }
120            Query::SetEdge { e, a, b } => {
121                dp.set_edge(e, (a, b));
122                pp!(dp.fold_all().sum);
123            }
124        }
125    }
126}
library_checker/tree/point_set_tree_path_composite_sum_fixed_root.rs

library_checker/tree/
point_set_tree_path_composite_sum_fixed_root.rs
