library_checker/tree/
point_set_tree_path_composite_sum.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
70#[derive(Clone, Copy, Debug, PartialEq, Eq)]
71struct PathPair {
72    forward: Path,
73    reverse: Path,
74}
75
76struct PathPairMonoid;
77impl Magma for PathPairMonoid {
78    type T = PathPair;
79    fn operate(x: &Self::T, y: &Self::T) -> Self::T {
80        PathPair {
81            forward: PathMonoid::operate(&x.forward, &y.forward),
82            reverse: PathMonoid::operate(&y.reverse, &x.reverse),
83        }
84    }
85}
86impl Unital for PathPairMonoid {
87    fn unit() -> Self::T {
88        PathPair {
89            forward: PathMonoid::unit(),
90            reverse: PathMonoid::unit(),
91        }
92    }
93}
94impl Associative for PathPairMonoid {}
95
96struct Dp;
97
98impl MonoidCluster for Dp {
99    type Vertex = MInt;
100    type Edge = (MInt, MInt);
101    type PointMonoid = PointMonoid;
102    type PathMonoid = PathPairMonoid;
103
104    fn add_vertex(point: &Point, vertex: &MInt, parent_edge: Option<&(MInt, MInt)>) -> PathPair {
105        let cnt = point.cnt + MInt::one();
106        let subtotal = point.sum + *vertex;
107        let (a, b) = parent_edge.copied().unwrap_or((MInt::one(), MInt::zero()));
108        PathPair {
109            forward: Path {
110                a,
111                b,
112                sum: a * subtotal + b * cnt,
113                cnt,
114            },
115            reverse: Path {
116                a,
117                b,
118                sum: subtotal,
119                cnt,
120            },
121        }
122    }
123
124    fn add_edge(path: &PathPair) -> Point {
125        Point {
126            sum: path.forward.sum,
127            cnt: path.forward.cnt,
128        }
129    }
130}
131
132competitive::define_enum_scan! {
133    enum Query: usize {
134        0 => SetVertex { v: usize, x: MInt, r: usize }
135        1 => SetEdge { e: usize, a: MInt, b: MInt, r: usize }
136    }
137}
138
139#[verify::library_checker("point_set_tree_path_composite_sum")]
140pub fn point_set_tree_path_composite_sum(reader: impl Read, mut writer: impl Write) {
141    let s = read_all_unchecked(reader);
142    let mut scanner = Scanner::new(&s);
143    scan!(
144        scanner,
145        n,
146        q,
147        value: [MInt; n],
148        (graph, edges): @TreeGraphScanner::<usize, (MInt, MInt)>::new(n)
149    );
150
151    let top_tree = graph.static_top_tree(0);
152    let mut dp = top_tree.dp::<Dp>(value, edges);
153
154    for _ in 0..q {
155        scan!(scanner, query: Query);
156        match query {
157            Query::SetVertex { v, x, r } => {
158                dp.set_vertex(v, x);
159                writeln!(writer, "{}", dp.fold_path(r).reverse.sum).ok();
160            }
161            Query::SetEdge { e, a, b, r } => {
162                dp.set_edge(e, (a, b));
163                writeln!(writer, "{}", dp.fold_path(r).reverse.sum).ok();
164            }
165        }
166    }
167}
library_checker/tree/point_set_tree_path_composite_sum.rs

library_checker/tree/
point_set_tree_path_composite_sum.rs
