haar_lib/ds/
palindromic_tree.rs1#![allow(clippy::len_without_is_empty)]
7
8use std::collections::{btree_map::Entry, BTreeMap};
9
10const ODD: usize = 0;
11const EVEN: usize = 1;
12
13#[derive(Default)]
15pub struct Node {
16 length: isize,
17 count: usize,
18 index: usize,
19 children: BTreeMap<char, *mut Node>,
20 parent: Option<*mut Node>,
21 suffix_link: Option<*mut Node>,
22 reverse_suffix_links: Vec<*mut Node>,
23}
24
25impl Node {
26 pub fn length(&self) -> isize {
28 self.length
29 }
30 pub fn count(&self) -> usize {
32 self.count
33 }
34 pub fn index(&self) -> usize {
36 self.index
37 }
38
39 pub fn parent(&self) -> Option<&Node> {
41 self.parent.map(|p| {
42 assert!(!p.is_null());
43 unsafe { &*p }
44 })
45 }
46
47 pub fn suffix_link(&self) -> Option<&Node> {
49 self.suffix_link.map(|p| {
50 assert!(!p.is_null());
51 unsafe { &*p }
52 })
53 }
54
55 pub fn children(&self) -> impl Iterator<Item = (char, &Node)> {
57 self.children.iter().map(|(&k, &v)| {
58 assert!(!v.is_null());
59 (k, unsafe { &*v })
60 })
61 }
62
63 pub fn rev_suffix_links(&self) -> impl Iterator<Item = &Node> {
65 self.reverse_suffix_links.iter().map(|&p| {
66 assert!(!p.is_null());
67 unsafe { &*p }
68 })
69 }
70}
71
72fn index_of(p: *mut Node) -> usize {
73 assert!(!p.is_null());
74 unsafe { (*p).index }
75}
76
77fn length_of(p: *mut Node) -> isize {
78 assert!(!p.is_null());
79 unsafe { (*p).length }
80}
81
82fn suffix_link_of(p: *mut Node) -> Option<*mut Node> {
83 assert!(!p.is_null());
84 unsafe { (*p).suffix_link }
85}
86
87fn child_of(p: *mut Node, c: char) -> Option<*mut Node> {
88 assert!(!p.is_null());
89 unsafe { (*p).children.get(&c).copied() }
90}
91
92fn set_suffix_link(from: *mut Node, to: *mut Node) {
93 assert!(!from.is_null());
94 assert!(!to.is_null());
95
96 unsafe {
97 (*from).suffix_link = Some(to);
98 (*to).reverse_suffix_links.push(from);
99 }
100}
101
102pub struct PalindromicTree {
104 list: Vec<*mut Node>,
105 sindex_list: Vec<usize>,
106 cur: *mut Node,
107 s: Vec<char>,
108}
109
110impl PalindromicTree {
111 pub fn new(s: &str) -> Self {
113 let even_root = Box::new(Node {
114 length: 0,
115 index: EVEN,
116 ..Default::default()
117 });
118 let even_root = Box::into_raw(even_root);
119
120 let odd_root = Box::new(Node {
121 length: -1,
122 index: ODD,
123 ..Default::default()
124 });
125 let odd_root = Box::into_raw(odd_root);
126
127 set_suffix_link(even_root, odd_root);
128
129 let s = s.chars().collect::<Vec<_>>();
130 let mut this = Self {
131 list: vec![odd_root, even_root],
132 sindex_list: vec![],
133 cur: odd_root,
134 s: vec![],
135 };
136
137 for c in s.into_iter() {
138 this.push(c);
139 }
140
141 this
142 }
143
144 pub fn push(&mut self, c: char) {
146 let i = self.s.len();
147 self.s.push(c);
148
149 let mut t = self.cur;
150
151 loop {
152 let t_index = index_of(t);
153
154 let k = i as isize - length_of(self.list[t_index]) - 1;
155 if k >= 0 && c == self.s[k as usize] {
156 let index = self.list.len();
157 let t = self.list[t_index];
158
159 assert!(!t.is_null());
160 match unsafe { (*t).children.entry(c) } {
161 Entry::Vacant(e) => {
162 let a = Box::new(Node {
163 length: length_of(t) + 2,
164 count: 1,
165 index,
166 parent: Some(self.list[index_of(t)]),
167 ..Default::default()
168 });
169
170 self.sindex_list.push(a.index);
171
172 let a = Box::into_raw(a);
173
174 e.insert(a);
175 self.list.push(a);
176 }
177 Entry::Occupied(e) => {
178 let t = index_of(*e.get());
179
180 assert!(!self.list[t].is_null());
181 unsafe { (*self.list[t]).count += 1 };
182 self.sindex_list.push(index_of(self.list[t]));
183 }
184 }
185
186 break;
187 }
188
189 t = suffix_link_of(self.list[t_index]).unwrap();
190 }
191
192 let next = child_of(self.list[index_of(t)], c).unwrap();
193 let next_index = index_of(next);
194
195 if suffix_link_of(self.list[next_index]).is_none() {
196 if length_of(self.list[next_index]) == 1 {
197 set_suffix_link(next, self.list[EVEN]);
198 } else {
199 let mut p = self.cur;
200
201 loop {
202 let p_index = index_of(p);
203
204 if p != t {
205 let k = i as isize - length_of(self.list[p_index]) - 1;
206
207 if k >= 0 && c == self.s[k as usize] {
208 let ch = child_of(self.list[p_index], c).unwrap();
209 set_suffix_link(next, ch);
210 break;
211 }
212 }
213 p = suffix_link_of(self.list[p_index]).unwrap();
214 }
215 }
216 }
217
218 self.cur = next;
219 }
220
221 pub fn len(&self) -> usize {
223 self.list.len()
224 }
225
226 pub fn odd_root(&self) -> &Node {
228 unsafe { &*self.list[ODD] }
229 }
230
231 pub fn even_root(&self) -> &Node {
233 unsafe { &*self.list[EVEN] }
234 }
235
236 pub fn node_of(&self, index: usize) -> Option<&Node> {
238 self.list.get(index).map(|&p| {
239 assert!(!p.is_null());
240 unsafe { &*p }
241 })
242 }
243
244 pub fn node_from_strpos(&self, pos: usize) -> Option<&Node> {
246 self.node_of(self.sindex_list[pos])
247 }
248}