Skip to main content

sa_is

Function sa_is 

Source
fn sa_is(s: &[usize], upper: usize) -> Vec<usize>
Examples found in repository?
crates/competitive/src/string/suffix_array.rs (line 25)
9    pub fn new<T>(text: &[T]) -> Self
10    where
11        T: Ord,
12    {
13        let n = text.len();
14        let mut ord: Vec<_> = (0..n).collect();
15        ord.sort_unstable_by_key(|&i| &text[i]);
16        let mut s = vec![0usize; n + 1];
17        let mut upper = 0usize;
18        for (k, &i) in ord.iter().enumerate() {
19            if k == 0 || text[ord[k - 1]] != text[i] {
20                upper += 1;
21            }
22            s[i] = upper;
23        }
24        s[n] = 0;
25        let sa = sa_is(&s, upper);
26        Self { sa }
27    }
28
29    pub fn into_inner(self) -> Vec<usize> {
30        self.sa
31    }
32
33    pub fn lcp_array_with_rank<T>(&self, text: &[T]) -> (Vec<usize>, Vec<usize>)
34    where
35        T: PartialEq,
36    {
37        let n = self.sa.len() - 1;
38        let mut rank = vec![0usize; n + 1];
39        for i in 0..=n {
40            rank[self.sa[i]] = i;
41        }
42
43        let mut h = 0usize;
44        let mut lcp_array = vec![0usize; n];
45        for i in 0..n {
46            let r = rank[i] - 1;
47            let j = self.sa[r];
48            while i + h < n && j + h < n && text[i + h] == text[j + h] {
49                h += 1;
50            }
51            lcp_array[r] = h;
52            h = h.saturating_sub(1);
53        }
54        (lcp_array, rank)
55    }
56
57    pub fn lcp_array<T>(&self, text: &[T]) -> Vec<usize>
58    where
59        T: PartialEq,
60    {
61        self.lcp_array_with_rank(text).0
62    }
63}
64
65impl Index<usize> for SuffixArray {
66    type Output = usize;
67    fn index(&self, index: usize) -> &Self::Output {
68        &self.sa[index]
69    }
70}
71
72fn induced_sort(s: &[usize], upper: usize, is_s: &[bool], lms: &[usize]) -> Vec<usize> {
73    let n = s.len();
74    let mut sa = vec![!0usize; n];
75
76    let mut bucket_size = vec![0usize; upper + 1];
77    for &c in s {
78        bucket_size[c] += 1;
79    }
80
81    let mut bucket_tail = vec![0usize; upper + 1];
82    {
83        let mut sum = 0usize;
84        for i in 0..=upper {
85            sum += bucket_size[i];
86            bucket_tail[i] = sum;
87        }
88    }
89    for &pos in lms.iter().rev() {
90        let c = s[pos];
91        bucket_tail[c] -= 1;
92        sa[bucket_tail[c]] = pos;
93    }
94
95    let mut bucket_head = vec![0usize; upper + 1];
96    {
97        let mut sum = 0usize;
98        for i in 0..=upper {
99            bucket_head[i] = sum;
100            sum += bucket_size[i];
101        }
102    }
103    for i in 0..n {
104        let v = sa[i];
105        if v == !0 || v == 0 {
106            continue;
107        }
108        let p = v - 1;
109        if !is_s[p] {
110            let c = s[p];
111            sa[bucket_head[c]] = p;
112            bucket_head[c] += 1;
113        }
114    }
115
116    let mut bucket_tail = vec![0usize; upper + 1];
117    {
118        let mut sum = 0usize;
119        for i in 0..=upper {
120            sum += bucket_size[i];
121            bucket_tail[i] = sum;
122        }
123    }
124    for i in (0..n).rev() {
125        let v = sa[i];
126        if v == !0 || v == 0 {
127            continue;
128        }
129        let p = v - 1;
130        if is_s[p] {
131            let c = s[p];
132            bucket_tail[c] -= 1;
133            sa[bucket_tail[c]] = p;
134        }
135    }
136
137    sa
138}
139
140fn sa_is(s: &[usize], upper: usize) -> Vec<usize> {
141    let n = s.len();
142    match n {
143        0 => return vec![],
144        1 => return vec![0],
145        _ => {}
146    }
147
148    let mut is_s = vec![false; n];
149    is_s[n - 1] = true;
150    for i in (0..n - 1).rev() {
151        is_s[i] = if s[i] < s[i + 1] {
152            true
153        } else if s[i] > s[i + 1] {
154            false
155        } else {
156            is_s[i + 1]
157        };
158    }
159
160    let mut lms_map = vec![!0; n];
161    let mut lms = vec![];
162    for i in 1..n {
163        if is_s[i] && !is_s[i - 1] {
164            lms_map[i] = lms.len();
165            lms.push(i);
166        }
167    }
168
169    let mut sa = induced_sort(s, upper, &is_s, &lms);
170
171    if lms.len() > 1 {
172        let mut rec_s = vec![0usize; lms.len()];
173        let mut rec_upper = 0usize;
174        let mut prev = !0usize;
175        for &pos in &sa {
176            if pos != !0 && lms_map[pos] != !0 {
177                if prev != !0 {
178                    let mut offset = 0usize;
179                    loop {
180                        if s[prev + offset] != s[pos + offset]
181                            || is_s[prev + offset] != is_s[pos + offset]
182                        {
183                            rec_upper += 1;
184                            break;
185                        }
186                        if offset != 0 {
187                            let prev_is_lms = lms_map[prev + offset] != !0;
188                            let pos_is_lms = lms_map[pos + offset] != !0;
189                            if prev_is_lms || pos_is_lms {
190                                rec_upper += (prev_is_lms != pos_is_lms) as usize;
191                                break;
192                            }
193                        }
194                        offset += 1;
195                    }
196                }
197                rec_s[lms_map[pos]] = rec_upper;
198                prev = pos;
199            }
200        }
201
202        let rec_sa = if rec_upper + 1 == lms.len() {
203            let mut rec_sa = vec![0usize; lms.len()];
204            for i in 0..lms.len() {
205                rec_sa[rec_s[i]] = i;
206            }
207            rec_sa
208        } else {
209            sa_is(&rec_s, rec_upper)
210        };
211
212        let mut ordered_lms = Vec::with_capacity(lms.len());
213        for &i in &rec_sa {
214            ordered_lms.push(lms[i]);
215        }
216        sa = induced_sort(s, upper, &is_s, &ordered_lms);
217    }
218
219    sa
220}