link_cut_tree.rs 11.8 KB
Newer Older
tgerstel's avatar
tgerstel committed
1
2
3
4
5
6
7
use super::collections::{Label, Pool};

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Node<V> {
    value: V,
    path_parent: Option<Label<Node<V>>>,
    parent: Option<Label<Node<V>>>,
SirBlueRabbit's avatar
SirBlueRabbit committed
8
    flipped: bool,
SirBlueRabbit's avatar
SirBlueRabbit committed
9
    children: [Option<Label<Node<V>>>; 2],
tgerstel's avatar
tgerstel committed
10
11
}

SirBlueRabbit's avatar
SirBlueRabbit committed
12
13
14
impl<V> Pool<Node<V>> {
    pub fn add_node(&mut self, value: V) -> Label<Node<V>> {
        self.insert(Node {
tgerstel's avatar
tgerstel committed
15
16
17
            value,
            path_parent: None,
            parent: None,
SirBlueRabbit's avatar
SirBlueRabbit committed
18
            flipped: false,
SirBlueRabbit's avatar
SirBlueRabbit committed
19
            children: [None, None],
SirBlueRabbit's avatar
SirBlueRabbit committed
20
        })
tgerstel's avatar
tgerstel committed
21
22
    }

SirBlueRabbit's avatar
SirBlueRabbit committed
23
    pub fn find_root(&mut self, label: Label<Node<V>>) -> Label<Node<V>> {
tgerstel's avatar
tgerstel committed
24
        self.expose(label);
tgerstel's avatar
tgerstel committed
25
        let mut root = label;
SirBlueRabbit's avatar
SirBlueRabbit committed
26
27
        let mut flipped = self[label].flipped;
        while let Some(next) = self[root].children[usize::from(flipped)] {
tgerstel's avatar
tgerstel committed
28
            root = next;
SirBlueRabbit's avatar
SirBlueRabbit committed
29
            flipped ^= self[next].flipped;
tgerstel's avatar
tgerstel committed
30
        }
tgerstel's avatar
tgerstel committed
31
32
33
34
35
36
        self.expose(root);
        root
    }

    pub fn cut(&mut self, label: Label<Node<V>>) {
        self.expose(label);
SirBlueRabbit's avatar
SirBlueRabbit committed
37
38
        let flipped = self[label].flipped;
        if let Some(left_label) = self[label].children[usize::from(flipped)] {
tgerstel's avatar
tgerstel committed
39
            self[left_label].parent = None;
SirBlueRabbit's avatar
SirBlueRabbit committed
40
            self[label].children[usize::from(flipped)] = None;
tgerstel's avatar
tgerstel committed
41
        }
tgerstel's avatar
tgerstel committed
42
43
44
45
46
    }

    pub fn link(&mut self, parent_label: Label<Node<V>>, child_label: Label<Node<V>>) {
        self.expose(child_label);
        self.expose(parent_label);
SirBlueRabbit's avatar
SirBlueRabbit committed
47
48
49
50
51
52
53
54
55
        let child_flipped = self[child_label].flipped;
        self[child_label].children[usize::from(child_flipped)] = Some(parent_label);
        self[parent_label].parent = Some(child_label);
    }

    pub fn evert(&mut self, label: Label<Node<V>>) {
        self.expose(label);
        self[label].flipped = !self[label].flipped;
    }
tgerstel's avatar
tgerstel committed
56
57
58
59
60
61
62
63
64

    pub fn value_mut(&mut self, label: Label<Node<V>>) -> &mut V {
        &mut self[label].value
    }

    pub fn value(&self, label: Label<Node<V>>) -> &V {
        &self[label].value
    }

SirBlueRabbit's avatar
SirBlueRabbit committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    fn expose(&mut self, v: Label<Node<V>>) {
        self.splay(v);
        let v_not_flipped = !self[v].flipped;
        if let Some(r) = self[v].children[usize::from(v_not_flipped)] {
            self[r].path_parent = Some(v);
            self[r].parent = None;
            self[v].children[usize::from(v_not_flipped)] = None;
        }

        while let Some(p) = self[v].path_parent {
            self.splay(p);
            let p_not_flipped = !self[p].flipped;
            if let Some(r) = self[p].children[usize::from(p_not_flipped)] {
                self[r].path_parent = Some(p);
                self[r].parent = None;
            }
            self[p].children[usize::from(p_not_flipped)] = Some(v);
            self[v].parent = Some(p);
            self[v].path_parent = None;
            self.rotate(v);
tgerstel's avatar
tgerstel committed
85
        }
SirBlueRabbit's avatar
SirBlueRabbit committed
86
    }
SirBlueRabbit's avatar
SirBlueRabbit committed
87
88
89
90

    pub fn splay(&mut self, label: Label<Node<V>>) {
        while let Some(p_label) = self[label].parent {
            if self[p_label].parent.is_some() {
SirBlueRabbit's avatar
SirBlueRabbit committed
91
                if self[p_label].flipped
SirBlueRabbit's avatar
SirBlueRabbit committed
92
93
                    == (self.is_zero_child(label) == self.is_zero_child(p_label))
                {
SirBlueRabbit's avatar
SirBlueRabbit committed
94
95
                    self.rotate(p_label);
                    self.rotate(label);
tgerstel's avatar
tgerstel committed
96
                } else {
SirBlueRabbit's avatar
SirBlueRabbit committed
97
98
                    self.rotate(label);
                    self.rotate(label);
tgerstel's avatar
tgerstel committed
99
100
                };
            } else {
SirBlueRabbit's avatar
SirBlueRabbit committed
101
                self.rotate(label);
tgerstel's avatar
tgerstel committed
102
103
                break;
            }
tgerstel's avatar
tgerstel committed
104
105
106
        }
    }

SirBlueRabbit's avatar
SirBlueRabbit committed
107
    fn rotate(&mut self, label: Label<Node<V>>) {
SirBlueRabbit's avatar
SirBlueRabbit committed
108
        let flipped = self[label].flipped;
SirBlueRabbit's avatar
SirBlueRabbit committed
109
        let p_side = !self.is_zero_child(label);
SirBlueRabbit's avatar
SirBlueRabbit committed
110
        let p_label = self[label].parent.unwrap();
SirBlueRabbit's avatar
SirBlueRabbit committed
111
        let b_option = self[label].children[usize::from(flipped ^ !p_side)];
SirBlueRabbit's avatar
SirBlueRabbit committed
112
113
114
115
        let g_option = self[p_label].parent;

        self[label].parent = g_option;
        if let Some(g_label) = g_option {
SirBlueRabbit's avatar
SirBlueRabbit committed
116
            let g_side = !self.is_zero_child(p_label);
SirBlueRabbit's avatar
SirBlueRabbit committed
117
118
119
120
121
            self[g_label].children[usize::from(g_side)] = Some(label);
        } else {
            self[label].path_parent = self[p_label].path_parent;
            self[p_label].path_parent = None;
        }
SirBlueRabbit's avatar
SirBlueRabbit committed
122
        self[p_label].children[usize::from(p_side)] = b_option;
SirBlueRabbit's avatar
SirBlueRabbit committed
123
        self[label].children[usize::from(flipped ^ !p_side)] = Some(p_label);
SirBlueRabbit's avatar
SirBlueRabbit committed
124
125
126
127
128

        if let Some(b_label) = b_option {
            self[b_label].parent = Some(p_label);
        }
        self[p_label].parent = Some(label);
SirBlueRabbit's avatar
SirBlueRabbit committed
129
130
131
132
133
134
    }

    fn is_zero_child(&self, label: Label<Node<V>>) -> bool {
        let parent_label = self[label].parent.unwrap();
        if let Some(zero_child_label) = self[parent_label].children[0] {
            zero_child_label == label
tgerstel's avatar
tgerstel committed
135
136
137
138
139
        } else {
            false
        }
    }
}
SirBlueRabbit's avatar
SirBlueRabbit committed
140
141
142
143
144
145
146
147

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_works() {
        let mut lct = Pool::with_capacity(1);
SirBlueRabbit's avatar
SirBlueRabbit committed
148
149
150
        let v = lct.add_node(());
        lct.splay(v);
        lct.expose(v);
SirBlueRabbit's avatar
SirBlueRabbit committed
151
152
153
154
155
156
157
    }

    #[test]
    fn rotate_basic() {
        let mut lct = Pool::with_capacity(2);

        // define nodes
SirBlueRabbit's avatar
SirBlueRabbit committed
158
159
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
SirBlueRabbit's avatar
SirBlueRabbit committed
160
161

        // link nodes
SirBlueRabbit's avatar
SirBlueRabbit committed
162
163
        lct[p].children[0] = Some(o);
        lct[o].parent = Some(p);
SirBlueRabbit's avatar
SirBlueRabbit committed
164
165

        // rotate
SirBlueRabbit's avatar
SirBlueRabbit committed
166
167
        lct.rotate(o);
        dbg!(&lct);
SirBlueRabbit's avatar
SirBlueRabbit committed
168
169

        // check children
SirBlueRabbit's avatar
SirBlueRabbit committed
170
171
        assert_eq!(lct[o].children, [None, Some(p)]);
        assert_eq!(lct[p].children, [None, None]);
SirBlueRabbit's avatar
SirBlueRabbit committed
172
173

        // check parents
SirBlueRabbit's avatar
SirBlueRabbit committed
174
175
        assert_eq!(lct[o].parent, None);
        assert_eq!(lct[p].parent, Some(o));
SirBlueRabbit's avatar
SirBlueRabbit committed
176
177
178
179
180
181
182
    }

    #[test]
    fn rotate_attachments() {
        let mut lct = Pool::with_capacity(7);

        // define nodes
SirBlueRabbit's avatar
SirBlueRabbit committed
183
184
185
186
187
188
189
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");
SirBlueRabbit's avatar
SirBlueRabbit committed
190
191

        // link nodes
SirBlueRabbit's avatar
SirBlueRabbit committed
192
193
194
195
196
197
198
199
200
        lct[o].children = [Some(a), Some(b)];
        lct[p].children = [Some(o), Some(c)];
        lct[g].children = [Some(p), Some(d)];
        lct[o].parent = Some(p);
        lct[p].parent = Some(g);
        lct[a].parent = Some(o);
        lct[b].parent = Some(o);
        lct[c].parent = Some(p);
        lct[d].parent = Some(g);
SirBlueRabbit's avatar
SirBlueRabbit committed
201
202

        // rotate
SirBlueRabbit's avatar
SirBlueRabbit committed
203
204
        lct.rotate(o);
        dbg!(&lct);
SirBlueRabbit's avatar
SirBlueRabbit committed
205
206

        // check children
SirBlueRabbit's avatar
SirBlueRabbit committed
207
208
209
210
211
212
213
        assert_eq!(lct[o].children, [Some(a), Some(p)]);
        assert_eq!(lct[p].children, [Some(b), Some(c)]);
        assert_eq!(lct[g].children, [Some(o), Some(d)]);
        assert_eq!(lct[a].children, [None, None]);
        assert_eq!(lct[b].children, [None, None]);
        assert_eq!(lct[c].children, [None, None]);
        assert_eq!(lct[d].children, [None, None]);
SirBlueRabbit's avatar
SirBlueRabbit committed
214
215

        // check parents
SirBlueRabbit's avatar
SirBlueRabbit committed
216
217
218
219
220
221
222
        assert_eq!(lct[o].parent, Some(g));
        assert_eq!(lct[p].parent, Some(o));
        assert_eq!(lct[g].parent, None);
        assert_eq!(lct[a].parent, Some(o));
        assert_eq!(lct[b].parent, Some(p));
        assert_eq!(lct[c].parent, Some(p));
        assert_eq!(lct[d].parent, Some(g));
SirBlueRabbit's avatar
SirBlueRabbit committed
223
    }
SirBlueRabbit's avatar
SirBlueRabbit committed
224
225

    #[test]
SirBlueRabbit's avatar
SirBlueRabbit committed
226
    fn rotate_flipped() {
SirBlueRabbit's avatar
SirBlueRabbit committed
227
228
229
        let mut lct = Pool::with_capacity(7);

        // define nodes
SirBlueRabbit's avatar
SirBlueRabbit committed
230
231
232
233
234
235
236
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");
SirBlueRabbit's avatar
SirBlueRabbit committed
237
238

        // link nodes
SirBlueRabbit's avatar
SirBlueRabbit committed
239
240
241
242
243
244
245
246
247
        lct[o].children = [Some(a), Some(b)];
        lct[p].children = [Some(o), Some(c)];
        lct[g].children = [Some(p), Some(d)];
        lct[o].parent = Some(p);
        lct[p].parent = Some(g);
        lct[a].parent = Some(o);
        lct[b].parent = Some(o);
        lct[c].parent = Some(p);
        lct[d].parent = Some(g);
SirBlueRabbit's avatar
SirBlueRabbit committed
248
249

        // reverse origin
SirBlueRabbit's avatar
SirBlueRabbit committed
250
        lct[o].flipped = true;
SirBlueRabbit's avatar
SirBlueRabbit committed
251
252

        // rotate
SirBlueRabbit's avatar
SirBlueRabbit committed
253
254
        lct.rotate(o);
        dbg!(&lct);
SirBlueRabbit's avatar
SirBlueRabbit committed
255
256

        // check children
SirBlueRabbit's avatar
SirBlueRabbit committed
257
258
259
260
261
262
263
        assert_eq!(lct[o].children, [Some(p), Some(b)]);
        assert_eq!(lct[p].children, [Some(a), Some(c)]);
        assert_eq!(lct[g].children, [Some(o), Some(d)]);
        assert_eq!(lct[a].children, [None, None]);
        assert_eq!(lct[b].children, [None, None]);
        assert_eq!(lct[c].children, [None, None]);
        assert_eq!(lct[d].children, [None, None]);
SirBlueRabbit's avatar
SirBlueRabbit committed
264
265

        // check parents
SirBlueRabbit's avatar
SirBlueRabbit committed
266
267
268
269
270
271
272
        assert_eq!(lct[o].parent, Some(g));
        assert_eq!(lct[p].parent, Some(o));
        assert_eq!(lct[g].parent, None);
        assert_eq!(lct[a].parent, Some(p));
        assert_eq!(lct[b].parent, Some(o));
        assert_eq!(lct[c].parent, Some(p));
        assert_eq!(lct[d].parent, Some(g));
tgerstel's avatar
tgerstel committed
273
274
275
    }

    #[test]
SirBlueRabbit's avatar
SirBlueRabbit committed
276
277
    fn expose_line() {
        let mut lct = Pool::with_capacity(5);
tgerstel's avatar
tgerstel committed
278

SirBlueRabbit's avatar
SirBlueRabbit committed
279
280
281
282
283
284
        // define nodes
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");
        let e = lct.add_node("e");
tgerstel's avatar
tgerstel committed
285

SirBlueRabbit's avatar
SirBlueRabbit committed
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
        // link nodes
        lct[b].children = [None, Some(c)];
        lct[d].children = [None, Some(e)];
        lct[c].parent = Some(b);
        lct[e].parent = Some(d);
        lct[b].path_parent = Some(a);
        lct[d].path_parent = Some(c);

        // expose
        lct.expose(d);
        dbg!(&lct);

        // check path parents
        assert_eq!(lct[a].path_parent, None);
        assert_eq!(lct[b].path_parent, None);
        assert_eq!(lct[c].path_parent, None);
        assert_eq!(lct[d].path_parent, None);
        assert_eq!(lct[e].path_parent, Some(d));
tgerstel's avatar
tgerstel committed
304
305
306
    }

    #[test]
SirBlueRabbit's avatar
SirBlueRabbit committed
307
308
    fn find_root() {
        let mut lct = Pool::with_capacity(7);
tgerstel's avatar
tgerstel committed
309

SirBlueRabbit's avatar
SirBlueRabbit committed
310
311
312
313
314
315
316
317
        // define nodes
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");
tgerstel's avatar
tgerstel committed
318

SirBlueRabbit's avatar
SirBlueRabbit committed
319
320
321
322
323
324
325
326
327
328
329
330
331
        // link nodes
        lct[o].children = [Some(a), Some(b)];
        lct[p].children = [Some(o), Some(c)];
        lct[g].children = [Some(p), Some(d)];
        lct[o].parent = Some(p);
        lct[p].parent = Some(g);
        lct[a].parent = Some(o);
        lct[b].parent = Some(o);
        lct[c].parent = Some(p);
        lct[d].parent = Some(g);

        // find root
        assert_eq!(lct.find_root(g), a);
tgerstel's avatar
tgerstel committed
332
333
334
    }

    #[test]
SirBlueRabbit's avatar
SirBlueRabbit committed
335
336
    fn cut_link() {
        let mut lct = Pool::with_capacity(5);
tgerstel's avatar
tgerstel committed
337

SirBlueRabbit's avatar
SirBlueRabbit committed
338
339
340
341
342
343
        // define nodes
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");
        let e = lct.add_node("e");
tgerstel's avatar
tgerstel committed
344

SirBlueRabbit's avatar
SirBlueRabbit committed
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
        // link nodes
        lct[b].children = [None, Some(c)];
        lct[d].children = [None, Some(e)];
        lct[c].parent = Some(b);
        lct[e].parent = Some(d);
        lct[b].path_parent = Some(a);
        lct[d].path_parent = Some(c);

        // cut
        lct.cut(c);
        dbg!(&lct);
        assert_eq!(lct[c].children[0], None);
        assert_eq!(lct[c].path_parent, None);
        assert_eq!(lct[b].children[1], None);

        // link
        lct.link(b, c);
        dbg!(&lct);
        assert_eq!(lct[b].parent, Some(c));
        assert_eq!(lct[c].children, [Some(b), None]);
        assert_eq!(lct[c].parent, None);
        assert_eq!(lct[d].path_parent, Some(c));
tgerstel's avatar
tgerstel committed
367
368
369
    }

    #[test]
SirBlueRabbit's avatar
SirBlueRabbit committed
370
371
    fn evert_root() {
        let mut lct = Pool::with_capacity(7);
tgerstel's avatar
tgerstel committed
372

SirBlueRabbit's avatar
SirBlueRabbit committed
373
374
375
376
377
378
379
380
        // define nodes
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");
tgerstel's avatar
tgerstel committed
381

SirBlueRabbit's avatar
SirBlueRabbit committed
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
        // link nodes
        lct[o].children = [Some(a), Some(b)];
        lct[p].children = [Some(o), Some(c)];
        lct[g].children = [Some(p), Some(d)];
        lct[o].parent = Some(p);
        lct[p].parent = Some(g);
        lct[a].parent = Some(o);
        lct[b].parent = Some(o);
        lct[c].parent = Some(p);
        lct[d].parent = Some(g);

        // evert
        lct.evert(o);
        dbg!(&lct);
        lct.expose(g);
        dbg!(&lct);
        assert_eq!(lct.find_root(g), o);
tgerstel's avatar
tgerstel committed
399
400
    }
}