Skip to main content

competitive/tree/
tree_hash.rs

1use crate::{
2    graph::{Graph, UndirectedSparseGraph},
3    tools::Xorshift,
4    tree::TreeCenter,
5};
6
7#[codesnip::entry("tree_hash", include("Xorshift", "tree_center"))]
8#[derive(Default, Debug)]
9pub struct TreeHasher {
10    rv: Vec<u64>,
11    rng: Xorshift,
12}
13#[codesnip::entry("tree_hash")]
14impl TreeHasher {
15    const MASK30: u64 = (1 << 30) - 1;
16    const MASK31: u64 = (1 << 31) - 1;
17    const MASK61: u64 = (1 << 61) - 1;
18    const MOD: u64 = Self::MASK61;
19    #[inline]
20    fn mersenne_mod(a: u64) -> u64 {
21        let mut res = (a >> 61) + (a & Self::MASK61);
22        if res >= Self::MASK61 {
23            res -= Self::MASK61;
24        }
25        res
26    }
27    #[inline]
28    fn mersenne_mul(a: u64, b: u64) -> u64 {
29        let au = a >> 31;
30        let ad = a & Self::MASK31;
31        let bu = b >> 31;
32        let bd = b & Self::MASK31;
33        let mid = ad * bu + au * bd;
34        let midu = mid >> 30;
35        let midd = mid & Self::MASK30;
36        au * bu * 2 + midu + (midd << 31) + ad * bd
37    }
38    #[inline]
39    fn mersenne_mul_mod(a: u64, b: u64) -> u64 {
40        Self::mersenne_mod(Self::mersenne_mul(a, b))
41    }
42    pub fn new() -> Self {
43        Self {
44            rv: Vec::new(),
45            rng: Xorshift::new(),
46        }
47    }
48    pub fn with_seed(seed: u64) -> Self {
49        Self {
50            rv: Vec::new(),
51            rng: Xorshift::new_with_seed(seed),
52        }
53    }
54    pub fn hash(&mut self, g: &UndirectedSparseGraph) -> u64 {
55        match g.tree_center() {
56            TreeCenter::One(u) => self.hash_rec(g, u, !0, 0),
57            TreeCenter::Two(u, v) => {
58                Self::mersenne_mul_mod(self.hash_rooted(g, u, v), self.hash_rooted(g, v, u))
59            }
60        }
61    }
62    pub fn hash_rooted(&mut self, g: &UndirectedSparseGraph, root: usize, parent: usize) -> u64 {
63        self.hash_rec(g, root, parent, 0)
64    }
65    fn hash_rec(&mut self, g: &UndirectedSparseGraph, u: usize, p: usize, d: usize) -> u64 {
66        let mut s = 1u64;
67        if self.rv.len() <= d {
68            self.rv.push(Self::mersenne_mod(self.rng.rand64()));
69        }
70        for a in g.neighbors(u) {
71            if a.to != p {
72                s = Self::mersenne_mul_mod(s, self.hash_rec(g, a.to, u, d + 1));
73            }
74        }
75        s += self.rv[d];
76        if s >= Self::MOD {
77            s -= Self::MOD;
78        }
79        s
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::tree::MixedTree;
87    use std::{
88        cmp::Ordering,
89        collections::{BTreeMap, HashMap},
90    };
91
92    fn vec_len_cmp<T: Ord>(left: &Vec<T>, right: &Vec<T>) -> Ordering {
93        match left.len().cmp(&right.len()) {
94            Ordering::Equal => left.cmp(right),
95            non_eq => non_eq,
96        }
97    }
98
99    impl UndirectedSparseGraph {
100        fn canonical(&self) -> Vec<bool> {
101            match self.tree_center() {
102                TreeCenter::One(u) => self.canonical_dfs(u, !0),
103                TreeCenter::Two(u, v) => {
104                    let mut a = self.canonical_dfs(u, v);
105                    let mut b = self.canonical_dfs(v, u);
106                    match vec_len_cmp(&a, &b) {
107                        Ordering::Less | Ordering::Equal => {
108                            a.append(&mut b);
109                            a
110                        }
111                        Ordering::Greater => {
112                            b.append(&mut a);
113                            b
114                        }
115                    }
116                }
117            }
118        }
119        fn canonical_dfs(&self, u: usize, p: usize) -> Vec<bool> {
120            let mut v = vec![vec![false]];
121            for a in self.neighbors(u) {
122                if a.to != p {
123                    v.push(self.canonical_dfs(a.to, u));
124                }
125            }
126            v.sort_unstable_by(vec_len_cmp);
127            v.push(vec![true]);
128            v.into_iter().flatten().collect()
129        }
130    }
131
132    #[test]
133    fn test_tree_hash() {
134        const N: usize = 200;
135        const Q: usize = 1000;
136        let mut rng = Xorshift::default();
137        let mut hasher = TreeHasher::new();
138        let mut h2s = HashMap::<u64, Vec<bool>>::new();
139        let mut s2h = BTreeMap::new();
140        for g in rng.random_iter(MixedTree(1..=N)).take(Q) {
141            let h = hasher.hash(&g);
142            let s = g.canonical();
143            h2s.entry(h)
144                .and_modify(|v| assert_eq!(v, &s))
145                .or_insert_with(|| s.clone());
146            s2h.entry(s).and_modify(|v| assert_eq!(*v, h)).or_insert(h);
147            assert_eq!(h2s.len(), s2h.len());
148        }
149        let mut v: Vec<_> = s2h.values().collect();
150        v.sort();
151        v.dedup();
152        assert_eq!(v.len(), s2h.len());
153        let mut v: Vec<_> = h2s.values().collect();
154        v.sort();
155        v.dedup();
156        assert_eq!(v.len(), h2s.len());
157    }
158}