contest-algorithms 0.1.3

Common algorithms and data structures for programming contests
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

//! Solves [Water Tree](http://codeforces.com/contest/343/problem/D).
//! To make a self-contained file for contest submission, dump each desired
//! module's contents directly here instead of the use statements.
//! Also, replace io::Cursor with io::stdin as shown in scanner.rs.
extern crate contest_algorithms;
use contest_algorithms::arq_tree::{ArqTree, AssignSum};
use contest_algorithms::graph::Graph;
use contest_algorithms::scanner::Scanner;

const SAMPLE_INPUT: &str = "\
5
1 2
5 1
2 3
4 2
12
1 1
2 3
3 1
3 2
3 3
3 4
1 2
2 4
3 1
3 3
3 4
3 5
";
const SAMPLE_OUTPUT: &str = "\
0
0
0
1
0
1
0
1
";

fn dfs(
    graph: &Graph,
    u: usize,
    l: &mut [usize],
    r: &mut [usize],
    p: &mut [usize],
    time: &mut usize,
) {
    *time += 1;
    l[u] = *time;

    for (_, v) in graph.adj_list(u) {
        if l[v] == 0 {
            p[v] = l[u];
            dfs(graph, v, l, r, p, time);
        }
    }

    r[u] = *time;
}

#[test]
fn main() {
    use std::fmt::Write;
    let cursor = std::io::Cursor::new(SAMPLE_INPUT);
    let mut scan = Scanner::new(cursor);
    let mut out = String::new();
    /* To read/write with stdin/stdout instead:
        use std::io::{self, Write};
        let (stdin, stdout) = (io::stdin(), io::stdout());
        let mut scan = Scanner::new(stdin.lock());
        let mut out = io::BufWriter::new(stdout.lock());
    */

    let n = scan.token::<usize>();
    let mut tree = Graph::new(n, 2 * (n - 1));
    for _ in 1..n {
        let u = scan.token::<usize>() - 1;
        let v = scan.token::<usize>() - 1;
        tree.add_undirected_edge(u, v);
    }

    let mut l = vec![0; n];
    let mut r = vec![0; n];
    let mut p = vec![0; n];
    dfs(&tree, 0, &mut l, &mut r, &mut p, &mut 0);

    let mut arq = ArqTree::<AssignSum>::new(vec![(0, 1); n + 1]);
    let q = scan.token::<usize>();
    for _ in 0..q {
        let c = scan.token::<usize>();
        let v = scan.token::<usize>() - 1;
        let (sum, len) = arq.query(l[v], r[v]);
        if c == 1 {
            if sum != len {
                arq.modify(p[v], p[v], &0);
                arq.modify(l[v], r[v], &1);
            }
        } else if c == 2 {
            arq.modify(l[v], l[v], &0);
        } else {
            let ans = if sum == len { 1 } else { 0 };
            writeln!(out, "{}", ans).ok();
        }
    }

    assert_eq!(out, SAMPLE_OUTPUT);
}