competitive/data_structure/
trie.rs

1use super::Monoid;
2use std::collections::VecDeque;
3
4#[derive(Debug, Clone)]
5pub struct Trie {
6    child: Vec<usize>,
7    node_size: usize,
8    char_size: usize,
9}
10impl Trie {
11    pub fn new(char_size: usize) -> Self {
12        Self {
13            child: vec![!0; char_size],
14            node_size: 1,
15            char_size,
16        }
17    }
18    pub fn with_capacity(char_size: usize, capacity: usize) -> Self {
19        let mut child = Vec::with_capacity(capacity * char_size);
20        child.resize(char_size, !0);
21        Self {
22            child,
23            node_size: 1,
24            char_size,
25        }
26    }
27    pub fn insert_once_at(&mut self, node: usize, ch: usize) -> usize {
28        let index = node * self.char_size + ch;
29        if self.child[index] == !0 {
30            self.child[index] = self.node_size;
31            self.child.resize(self.child.len() + self.char_size, !0);
32            self.node_size += 1;
33        }
34        self.child[index]
35    }
36    pub fn insert_at<I>(&mut self, mut node: usize, iter: I) -> Vec<usize>
37    where
38        I: IntoIterator<Item = usize>,
39    {
40        let mut path = Vec::new();
41        for ch in iter.into_iter() {
42            path.push(node);
43            node = self.insert_once_at(node, ch);
44        }
45        path.push(node);
46        path
47    }
48    pub fn insert<I>(&mut self, iter: I) -> Vec<usize>
49    where
50        I: IntoIterator<Item = usize>,
51    {
52        self.insert_at(0, iter)
53    }
54    pub fn find_at<I>(&self, mut node: usize, iter: I) -> Result<usize, usize>
55    where
56        I: IntoIterator<Item = usize>,
57    {
58        for ch in iter.into_iter() {
59            if let Some(&nnode) = self.child.get(node * self.char_size + ch) {
60                node = nnode;
61            } else {
62                return Err(node);
63            }
64        }
65        Ok(node)
66    }
67    pub fn find<I>(&self, iter: I) -> Result<usize, usize>
68    where
69        I: IntoIterator<Item = usize>,
70    {
71        self.find_at(0, iter)
72    }
73    pub fn next_node(&self, node: usize, ch: usize) -> Option<usize> {
74        let index = node * self.char_size + ch;
75        if self.child[index] == !0 {
76            None
77        } else {
78            Some(self.child[index])
79        }
80    }
81    pub fn node_size(&self) -> usize {
82        self.node_size
83    }
84    pub fn edges(&self) -> Vec<(usize, usize)> {
85        let mut edges = Vec::with_capacity(self.node_size - 1);
86        let mut stack = vec![0usize];
87        while let Some(node) = stack.pop() {
88            for ch in (0..self.char_size).rev() {
89                if let Some(nnode) = self.next_node(node, ch) {
90                    edges.push((node, nnode));
91                    stack.push(nnode);
92                }
93            }
94        }
95        edges
96    }
97    pub fn build_failure<M>(&mut self, dp: &mut [M::T])
98    where
99        M: Monoid,
100    {
101        let mut fail = vec![0usize; self.node_size];
102        let mut deq = VecDeque::new();
103        for ch in 0..self.char_size {
104            if self.child[ch] != !0 {
105                deq.push_back(self.child[ch]);
106            } else {
107                self.child[ch] = 0;
108            }
109        }
110        while let Some(node) = deq.pop_front() {
111            let f = fail[node];
112            dp[node] = M::operate(&dp[node], &dp[f]);
113            let base = node * self.char_size;
114            let fbase = f * self.char_size;
115            for ch in 0..self.char_size {
116                let index = base + ch;
117                let nnode = self.child[index];
118                if nnode != !0 {
119                    deq.push_back(nnode);
120                    fail[nnode] = self.child[fbase + ch];
121                } else {
122                    self.child[index] = self.child[fbase + ch];
123                }
124            }
125        }
126    }
127}
competitive/data_structure/trie.rs

competitive/data_structure/
trie.rs
