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    fn vec_len_cmp<T: Ord>(left: &Vec<T>, right: &Vec<T>) -> Ordering {
89        match left.len().cmp(&right.len()) {
90            Ordering::Equal => left.cmp(right),
91            non_eq => non_eq,
92        }
93    }
94
95    impl UndirectedSparseGraph {
96        fn canonical(&self) -> Vec<bool> {
97            match self.tree_center() {
98                TreeCenter::One(u) => self.canonical_dfs(u, !0),
99                TreeCenter::Two(u, v) => {
100                    let mut a = self.canonical_dfs(u, v);
101                    let mut b = self.canonical_dfs(v, u);
102                    match vec_len_cmp(&a, &b) {
103                        Ordering::Less | Ordering::Equal => {
104                            a.append(&mut b);
105                            a
106                        }
107                        Ordering::Greater => {
108                            b.append(&mut a);
109                            b
110                        }
111                    }
112                }
113            }
114        }
115        fn canonical_dfs(&self, u: usize, p: usize) -> Vec<bool> {
116            let mut v = vec![vec![false]];
117            for a in self.adjacencies(u) {
118                if a.to != p {
119                    v.push(self.canonical_dfs(a.to, u));
120                }
121            }
122            v.sort_unstable_by(vec_len_cmp);
123            v.push(vec![true]);
124            v.into_iter().flatten().collect()
125        }
126    }
127
128    #[test]
129    fn test_tree_hash() {
130        const N: usize = 200;
131        const Q: usize = 1000;
132        let mut rng = Xorshift::default();
133        let mut hasher = TreeHasher::new();
134        let mut h2s = HashMap::<u64, Vec<bool>>::new();
135        let mut s2h = BTreeMap::new();
136        for g in rng.random_iter(MixedTree(1..=N)).take(Q) {
137            let h = hasher.hash(&g);
138            let s = g.canonical();
139            h2s.entry(h)
140                .and_modify(|v| assert_eq!(v, &s))
141                .or_insert_with(|| s.clone());
142            s2h.entry(s).and_modify(|v| assert_eq!(*v, h)).or_insert(h);
143            assert_eq!(h2s.len(), s2h.len());
144        }
145        let mut v: Vec<_> = s2h.values().collect();
146        v.sort();
147        v.dedup();
148        assert_eq!(v.len(), s2h.len());
149        let mut v: Vec<_> = h2s.values().collect();
150        v.sort();
151        v.dedup();
152        assert_eq!(v.len(), h2s.len());
153    }
154}