competitive/tree/
tree_hash.rs

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