haar_lib/tree/
rooting.rs

1//! 根付き木に変換
2
3use crate::tree::*;
4use std::mem::swap;
5
6/// 根付き木に変換
7pub fn rooting<E: TreeEdgeTrait>(tr: &mut Tree<E>, root: usize) -> Result<(), &str> {
8    let n = tr.len();
9    let mut stack = vec![(root, -1)];
10    let mut check = vec![false; n];
11
12    while let Some((cur, par)) = stack.pop() {
13        if check[cur] {
14            return Err("loop detected");
15        }
16        check[cur] = true;
17
18        if par == -1 {
19            if let Some(p) = tr.nodes[cur].parent.take() {
20                tr.nodes[cur].children.push(p);
21            }
22        } else if let Some(mut p) = tr.nodes[cur].parent.take() {
23            for e in tr.nodes[cur].children.iter_mut() {
24                if e.to() == par as usize {
25                    swap(&mut p, e);
26                    tr.nodes[cur].parent = Some(p);
27                    break;
28                }
29            }
30        } else {
31            for (i, e) in tr.nodes[cur].children.iter().enumerate() {
32                if e.to() == par as usize {
33                    let x = tr.nodes[cur].children.swap_remove(i);
34                    tr.nodes[cur].parent = Some(x);
35                    break;
36                }
37            }
38        }
39
40        for e in &tr.nodes[cur].children {
41            stack.push((e.to(), cur as isize));
42        }
43    }
44
45    Ok(())
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn test() {
54        let mut builder = TreeBuilder::new(6);
55        builder.extend(
56            vec![(0, 1), (1, 2), (2, 3), (2, 4), (5, 1)]
57                .into_iter()
58                .map(|(u, v)| TreeEdge::new(u, v, (), ())),
59        );
60        let mut tr = builder.build();
61        assert_eq!(rooting(&mut tr, 0), Ok(()));
62        assert_eq!(
63            tr.nodes
64                .into_iter()
65                .map(|nd| nd.parent.map(|p| p.to))
66                .collect::<Vec<_>>(),
67            vec![None, Some(0), Some(1), Some(2), Some(2), Some(1)]
68        );
69
70        let mut builder = TreeBuilder::new(6);
71        builder.extend(
72            vec![(0, 1), (1, 2), (2, 3), (2, 1), (5, 1)]
73                .into_iter()
74                .map(|(u, v)| TreeEdge::new(u, v, (), ())),
75        );
76        let mut tr = builder.build();
77        assert!(rooting(&mut tr, 0).is_err());
78    }
79}
haar_lib/tree/rooting.rs

haar_lib/tree/
rooting.rs
