Skip to main content

library_checker/tree/
tree_path_composite_sum.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::AdditiveOperation, graph::TreeGraphScanner, num::mint_basic::MInt998244353 as M,
4    tree::ReRooting,
5};
6
7#[verify::library_checker("tree_path_composite_sum")]
8pub fn tree_path_composite_sum(reader: impl Read, writer: impl Write) {
9    prepare_io!(reader, writer);
10    sc!(n, values: [M; n], (graph, edges): @TreeGraphScanner::<usize, (M, M)>::new(n));
11    let dp = ReRooting::<(AdditiveOperation<M>, AdditiveOperation<i32>), _>::new_with_inverse(
12        &graph,
13        |&(sum, count), v, edge| {
14            let sum = sum + values[v];
15            let count = count + 1;
16            if let Some(edge) = edge {
17                let (a, b) = edges[edge];
18                (a * sum + b * M::new_unchecked(count as u32), count)
19            } else {
20                (sum, count)
21            }
22        },
23    );
24    pp!(@it dp.dp.iter().map(|value| value.0));
25}