treez 1.6.0

A collection of useful data structures.
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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

#[derive(Debug, Clone, Copy)]
pub struct S {
    pub id: usize,
    pub parent: usize,
    pub ssize: usize,
}

impl S {
    pub fn init(id: usize) -> Self {
        Self {
            id: id,
            parent: id,
            ssize: 1,
        }
    }
}

pub struct Dsu(pub Vec<S>);

impl Dsu {
    pub fn init(len: usize) -> Self {
        Dsu((0..len).map(|x| S::init(x)).collect::<Vec<_>>())
    }

    pub fn insert_new(&mut self, count: usize) -> usize {
        let l = self.0.len();
        self.0.push(S::init(l));
        l
    }

    pub fn insert_new_count(&mut self, count: usize) -> Vec<usize> {
        let mut ret = vec![];
        let l = self.0.len();
        for i in l..l + count {
            self.0.push(S::init(i));
            ret.push(i);
        }
        ret
    }

    pub fn compress_path(&mut self, mut idx: usize) -> usize {
        debug_assert!(idx < self.0.len());

        let mut buf = vec![];

        while self.0[idx].parent != self.0[idx].id {
            self.0[idx].ssize = 1;
            buf.push(idx);
            idx = self.0[idx].parent;
        }

        for i in buf {
            self.0[i].parent = idx;
        }

        idx
    }

    pub fn merge(&mut self, a: usize, b: usize) {
        if a != b {
            let a_parent = self.compress_path(a);
            let b_parent = self.compress_path(b);

            if a_parent != b_parent {
                // different sets check
                let (small, big) = if self.0[a_parent].ssize < self.0[b_parent].ssize {
                    (a_parent, b_parent)
                } else {
                    (b_parent, a_parent)
                };
                self.0[big].ssize += self.0[small].ssize;
                self.0[small].ssize = 1;
                self.0[small].parent = big;
            }
        }
    }

    pub fn get_set(&self, id: usize) -> Option<usize> {
        if id >= self.0.len() {
            None
        } else {
            Some(self.0[id].parent)
        }
    }

    pub fn get_sets_repr(&self) -> Vec<usize> {
        self.0
            .iter()
            .enumerate()
            .filter(|x| x.1.id == x.1.parent)
            .map(|x| x.0)
            .collect()
    }

    pub fn get_sets_len(&self) -> usize {
        self.0.iter().filter(|x| x.id == x.parent).count()
    }

    pub fn get_set_size(&self, id: usize) -> usize {
        if id > self.0.len() {
            0
        } else {
            self.0[id].ssize
        }
    }
}

#[test]
fn test_dsu() {
    let mut v = Dsu::init(10);

    //1, 3, 5, 7 ,9
    for i in 0..5 {
        let j = i * 2 + 1;
        v.merge(j, j - 1);
    }

    let ret = v.get_sets_repr();
    assert_eq!(ret.len(), 5);

    v.merge(5, 9);

    assert_eq!(v.get_sets_repr().len(), 4);

    v.merge(7, 2);

    assert_eq!(v.get_sets_repr().len(), 3);

    v.merge(4, 3);

    assert_eq!(v.get_sets_repr().len(), 2);

    v.merge(0, 5);

    assert_eq!(v.get_sets_repr().len(), 1);

    let _ids = v.insert_new_count(5);
    assert_eq!(v.get_sets_repr().len(), 6);

    let r = *v.get_sets_repr().iter().rev().nth(0).unwrap();
    for i in v.get_sets_repr() {
        v.merge(i, r);
    }
    assert_eq!(v.get_sets_len(), 1);
}