haar_lib/ds/
palindromic_tree.rs

1//! 回文木
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/eertree>
5
6#![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/// [`PalindromicTree`]のノード
14#[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    /// 回文の長さを返す。
27    pub fn length(&self) -> isize {
28        self.length
29    }
30    /// ノードが表す回文の出現回数を返す。
31    pub fn count(&self) -> usize {
32        self.count
33    }
34    /// ノードに割り当てられた番号を返す。
35    pub fn index(&self) -> usize {
36        self.index
37    }
38
39    /// 親ノードへの参照を返す。
40    pub fn parent(&self) -> Option<&Node> {
41        self.parent.map(|p| {
42            assert!(!p.is_null());
43            unsafe { &*p }
44        })
45    }
46
47    /// 接尾辞リンクの行き先ノードへの参照を返す。
48    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    /// 子ノードへの推移文字と子ノードへの参照へのイテレータを返す。
56    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    /// 接尾辞リンクを逆に辿ったノードへの参照へのイテレータを返す。
64    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
102/// 回文木
103pub 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    /// 文字列`s`から回文木を構築する。
112    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    /// 末尾に文字`c`を追加する。
145    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    /// 回文木に含まれるノードの個数を返す。(長さ`0`,`-1`のノードも含む。)
222    pub fn len(&self) -> usize {
223        self.list.len()
224    }
225
226    /// 奇数長回文の木の根への参照を返す。
227    pub fn odd_root(&self) -> &Node {
228        unsafe { &*self.list[ODD] }
229    }
230
231    /// 偶数長回文の木の根への参照を返す。
232    pub fn even_root(&self) -> &Node {
233        unsafe { &*self.list[EVEN] }
234    }
235
236    /// `index`番目のノードへの参照を返す。
237    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    /// 元の文字列の長さ`pos+1`の接頭辞の最大回文接尾辞に対応するノードへの参照を返す。
245    pub fn node_from_strpos(&self, pos: usize) -> Option<&Node> {
246        self.node_of(self.sindex_list[pos])
247    }
248}