library_checker/tree/
point_set_tree_path_composite_sum_fixed_root.rs

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

library_checker/tree/
point_set_tree_path_composite_sum_fixed_root.rs
