competitive/algorithm/
quotient_index.rs

1use std::{
2    iter::{ExactSizeIterator, FusedIterator},
3    ops::RangeInclusive,
4};
5
6/// sorted({ floor(n/k) | k in \[1, n\] })
7pub struct FloorQuotientIndex {
8    num: usize,
9    sqrt: usize,
10    pivot: usize,
11}
12
13impl FloorQuotientIndex {
14    pub fn new(num: usize) -> Self {
15        assert!(num > 0, "num must be positive");
16        let sqrt = num.isqrt();
17        let pivot = num / sqrt;
18        Self { num, sqrt, pivot }
19    }
20
21    pub fn values(&self) -> impl ExactSizeIterator<Item = usize> + FusedIterator + '_ {
22        let len = self.len();
23        (0..len).map(move |index| unsafe { self.index_to_value_unchecked(index) })
24    }
25
26    pub fn key_values(
27        &self,
28    ) -> impl ExactSizeIterator<Item = (RangeInclusive<usize>, usize)> + FusedIterator + '_ {
29        let len = self.len();
30        (0..len).map(move |index| unsafe {
31            (
32                self.index_to_key_unchecked(index),
33                self.index_to_value_unchecked(index),
34            )
35        })
36    }
37
38    #[allow(clippy::len_without_is_empty)]
39    pub fn len(&self) -> usize {
40        2 * self.sqrt - (self.pivot == self.sqrt) as usize
41    }
42
43    pub fn contains_key(&self, key: usize) -> bool {
44        (1..=self.num).contains(&key)
45    }
46
47    pub fn contains_value(&self, value: usize) -> bool {
48        if value == 0 || value > self.num {
49            return false;
50        }
51        if value == usize::MAX {
52            return self.num == usize::MAX;
53        }
54        self.num / value > self.num / (value + 1)
55    }
56
57    pub fn index_to_key(&self, index: usize) -> RangeInclusive<usize> {
58        assert!(index < self.len(), "index out of bounds");
59        unsafe { self.index_to_key_unchecked(index) }
60    }
61
62    pub fn index_to_value(&self, index: usize) -> usize {
63        assert!(index < self.len(), "index out of bounds");
64        unsafe { self.index_to_value_unchecked(index) }
65    }
66
67    pub fn value_to_index(&self, value: usize) -> usize {
68        assert!(self.contains_value(value), "value is not present");
69        unsafe { self.value_to_index_unchecked(value) }
70    }
71
72    pub fn value_to_key(&self, value: usize) -> RangeInclusive<usize> {
73        assert!(self.contains_value(value), "value is not present");
74        unsafe { self.value_to_key_unchecked(value) }
75    }
76
77    pub fn key_to_index(&self, key: usize) -> usize {
78        assert!(self.contains_key(key), "key out of bounds");
79        let value = self.key_to_value(key);
80        unsafe { self.value_to_index_unchecked(value) }
81    }
82
83    pub fn key_to_value(&self, key: usize) -> usize {
84        assert!(self.contains_key(key), "key out of bounds");
85        unsafe { self.key_to_value_unchecked(key) }
86    }
87
88    /// # Safety
89    /// `index` must satisfy `index < self.len()`.
90    pub unsafe fn index_to_key_unchecked(&self, index: usize) -> RangeInclusive<usize> {
91        if index < self.sqrt {
92            self.num / (index + 2) + 1..=self.num / (index + 1)
93        } else {
94            let key = self.len() - index;
95            key..=key
96        }
97    }
98
99    /// # Safety
100    /// `index` must satisfy `index < self.len()`.
101    pub unsafe fn index_to_value_unchecked(&self, index: usize) -> usize {
102        if index < self.sqrt {
103            index + 1
104        } else {
105            self.num / (self.len() - index)
106        }
107    }
108
109    /// # Safety
110    /// `value` must satisfy `self.contains_value(value)`.
111    pub unsafe fn value_to_index_unchecked(&self, value: usize) -> usize {
112        if value <= self.sqrt {
113            value - 1
114        } else {
115            self.len() - self.num / value
116        }
117    }
118
119    /// # Safety
120    /// `value` must satisfy `self.contains_value(value)`.
121    pub unsafe fn value_to_key_unchecked(&self, value: usize) -> RangeInclusive<usize> {
122        if value == usize::MAX {
123            return 1..=1;
124        }
125        let start = (self.num / (value + 1)) + 1;
126        let end = self.num / value;
127        start..=end
128    }
129
130    /// # Safety
131    /// `key` must satisfy `1 <= key && key <= self.num`.
132    pub unsafe fn key_to_index_unchecked(&self, key: usize) -> usize {
133        let len = self.len();
134        if key <= len - self.sqrt {
135            len - key
136        } else {
137            self.num / key - 1
138        }
139    }
140
141    /// # Safety
142    /// `key` must satisfy `1 <= key && key <= self.num`.
143    pub unsafe fn key_to_value_unchecked(&self, key: usize) -> usize {
144        self.num / key
145    }
146}
147
148/// sorted({ ceil(n/k) | k in \[1, n\] })
149pub struct CeilQuotientIndex {
150    num: usize,
151    sqrt: usize,
152    pivot: usize,
153}
154
155impl CeilQuotientIndex {
156    pub fn new(num: usize) -> Self {
157        assert!(num > 0, "num must be positive");
158        let sqrt = num.isqrt();
159        let pivot = num.div_ceil(sqrt);
160        Self { num, sqrt, pivot }
161    }
162
163    pub fn values(&self) -> impl ExactSizeIterator<Item = usize> + FusedIterator + '_ {
164        let len = self.len();
165        (0..len).map(move |index| unsafe { self.index_to_value_unchecked(index) })
166    }
167
168    pub fn key_values(
169        &self,
170    ) -> impl ExactSizeIterator<Item = (RangeInclusive<usize>, usize)> + FusedIterator + '_ {
171        let len = self.len();
172        (0..len).map(move |index| unsafe {
173            (
174                self.index_to_key_unchecked(index),
175                self.index_to_value_unchecked(index),
176            )
177        })
178    }
179
180    #[allow(clippy::len_without_is_empty)]
181    pub fn len(&self) -> usize {
182        self.sqrt + self.pivot - 1
183    }
184
185    pub fn contains_key(&self, key: usize) -> bool {
186        (1..=self.num).contains(&key)
187    }
188
189    pub fn contains_value(&self, value: usize) -> bool {
190        if value == 0 || value > self.num {
191            return false;
192        }
193        if value == 1 {
194            return true;
195        }
196        let start = self.num.div_ceil(value);
197        let end = (self.num - 1) / (value - 1);
198        start <= end
199    }
200
201    pub fn index_to_key(&self, index: usize) -> RangeInclusive<usize> {
202        assert!(index < self.len(), "index out of bounds");
203        unsafe { self.index_to_key_unchecked(index) }
204    }
205
206    pub fn index_to_value(&self, index: usize) -> usize {
207        assert!(index < self.len(), "index out of bounds");
208        unsafe { self.index_to_value_unchecked(index) }
209    }
210
211    pub fn value_to_index(&self, value: usize) -> usize {
212        assert!(self.contains_value(value), "value is not present");
213        unsafe { self.value_to_index_unchecked(value) }
214    }
215
216    pub fn value_to_key(&self, value: usize) -> RangeInclusive<usize> {
217        assert!(self.contains_value(value), "value is not present");
218        unsafe { self.value_to_key_unchecked(value) }
219    }
220
221    pub fn key_to_index(&self, key: usize) -> usize {
222        assert!(self.contains_key(key), "key out of bounds");
223        let value = self.key_to_value(key);
224        unsafe { self.value_to_index_unchecked(value) }
225    }
226
227    pub fn key_to_value(&self, key: usize) -> usize {
228        assert!(self.contains_key(key), "key out of bounds");
229        unsafe { self.key_to_value_unchecked(key) }
230    }
231
232    /// # Safety
233    /// `index` must satisfy `index < self.len()`.
234    pub unsafe fn index_to_key_unchecked(&self, index: usize) -> RangeInclusive<usize> {
235        if index < self.pivot {
236            if index == 0 {
237                return self.num..=self.num;
238            }
239            let start = self.num.div_ceil(index + 1);
240            let end = (self.num - 1) / index;
241            start..=end
242        } else {
243            let key = self.len() - index;
244            key..=key
245        }
246    }
247
248    /// # Safety
249    /// `index` must satisfy `index < self.len()`.
250    pub unsafe fn index_to_value_unchecked(&self, index: usize) -> usize {
251        if index < self.pivot {
252            index + 1
253        } else {
254            self.num.div_ceil(self.len() - index)
255        }
256    }
257
258    /// # Safety
259    /// `value` must satisfy `self.contains_value(value)`.
260    pub unsafe fn value_to_index_unchecked(&self, value: usize) -> usize {
261        if value <= self.pivot {
262            value - 1
263        } else {
264            self.len() - (self.num - 1) / (value - 1)
265        }
266    }
267
268    /// # Safety
269    /// `value` must satisfy `self.contains_value(value)`.
270    pub unsafe fn value_to_key_unchecked(&self, value: usize) -> RangeInclusive<usize> {
271        if value == 1 {
272            return self.num..=self.num;
273        }
274        let start = self.num.div_ceil(value);
275        let end = (self.num - 1) / (value - 1);
276        start..=end
277    }
278
279    /// # Safety
280    /// `key` must satisfy `1 <= key && key <= self.num`.
281    pub unsafe fn key_to_index_unchecked(&self, key: usize) -> usize {
282        if key < self.sqrt {
283            self.len() - key
284        } else {
285            self.num.div_ceil(key) - 1
286        }
287    }
288
289    /// # Safety
290    /// `key` must satisfy `1 <= key && key <= self.num`.
291    pub unsafe fn key_to_value_unchecked(&self, key: usize) -> usize {
292        self.num.div_ceil(key)
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299    use crate::tools::Xorshift;
300
301    #[test]
302    fn test_quotient_index() {
303        let mut rng = Xorshift::default();
304        for n in (1..=2000).chain(rng.random_iter(1..=200_000).take(100)) {
305            let qi = FloorQuotientIndex::new(n);
306            let mut expected_values: Vec<_> = (1..=n).rev().map(|key| n / key).collect();
307            expected_values.dedup();
308            assert_eq!(qi.len(), expected_values.len());
309            let values: Vec<_> = qi.values().collect();
310            assert_eq!(values, expected_values);
311            let expected_key_values: Vec<_> = (0..qi.len())
312                .map(|index| (qi.index_to_key(index), qi.index_to_value(index)))
313                .collect();
314            let key_values: Vec<_> = qi.key_values().collect();
315            assert_eq!(key_values, expected_key_values);
316            let mut count = 0;
317            let mut is_value = vec![false; n + 1];
318            for (index, &value) in values.iter().enumerate() {
319                assert_eq!(qi.index_to_value(index), value);
320                assert_eq!(qi.value_to_index(value), index);
321                assert_eq!(qi.index_to_key(index), qi.value_to_key(value));
322                for key in qi.index_to_key(index) {
323                    assert_eq!(qi.key_to_value(key), value);
324                    assert_eq!(qi.key_to_index(key), index);
325                    count += 1;
326                }
327                is_value[value] = true;
328            }
329            assert_eq!(count, n);
330            for (value, &present) in is_value.iter().enumerate() {
331                assert_eq!(qi.contains_value(value), present);
332            }
333            assert!(!qi.contains_key(!0));
334        }
335
336        for n in [usize::MAX]
337            .into_iter()
338            .chain(rng.random_iter(1..=!0).take(100))
339        {
340            let qi = FloorQuotientIndex::new(n);
341            let len = qi.len();
342            for index in (0..len.min(1000)).chain(len.saturating_sub(1000)..len) {
343                let value = qi.index_to_value(index);
344                assert_eq!(qi.index_to_value(index), value);
345                assert_eq!(qi.value_to_index(value), index);
346                assert_eq!(qi.index_to_key(index), qi.value_to_key(value));
347                for key in qi
348                    .index_to_key(index)
349                    .take(1)
350                    .chain(qi.index_to_key(index).rev().take(1))
351                {
352                    assert_eq!(qi.key_to_value(key), value);
353                    assert_eq!(qi.key_to_index(key), index);
354                }
355            }
356        }
357    }
358
359    #[test]
360    fn test_ceil_quotient_index() {
361        let mut rng = Xorshift::default();
362        for n in (1..=2000).chain(rng.random_iter(1..=200_000).take(100)) {
363            let qi = CeilQuotientIndex::new(n);
364            let mut expected_values: Vec<_> = (1..=n).rev().map(|key| n.div_ceil(key)).collect();
365            expected_values.dedup();
366            assert_eq!(qi.len(), expected_values.len());
367            let values: Vec<_> = qi.values().collect();
368            assert_eq!(values, expected_values);
369            let expected_key_values: Vec<_> = (0..qi.len())
370                .map(|index| (qi.index_to_key(index), qi.index_to_value(index)))
371                .collect();
372            let key_values: Vec<_> = qi.key_values().collect();
373            assert_eq!(key_values, expected_key_values);
374            let mut count = 0;
375            let mut is_value = vec![false; n + 1];
376            for (index, &value) in values.iter().enumerate() {
377                assert_eq!(qi.index_to_value(index), value);
378                assert_eq!(qi.value_to_index(value), index);
379                assert_eq!(qi.index_to_key(index), qi.value_to_key(value));
380                for key in qi.index_to_key(index) {
381                    assert_eq!(qi.key_to_value(key), value);
382                    assert_eq!(qi.key_to_index(key), index);
383                    count += 1;
384                }
385                is_value[value] = true;
386            }
387            assert_eq!(count, n);
388            for (value, &present) in is_value.iter().enumerate() {
389                assert_eq!(qi.contains_value(value), present);
390            }
391            assert!(!qi.contains_key(!0));
392        }
393
394        for n in [usize::MAX]
395            .into_iter()
396            .chain(rng.random_iter(1..=!0).take(100))
397        {
398            let qi = CeilQuotientIndex::new(n);
399            let len = qi.len();
400            for index in (0..len.min(1000)).chain(len.saturating_sub(1000)..len) {
401                let value = qi.index_to_value(index);
402                assert_eq!(qi.index_to_value(index), value);
403                assert_eq!(qi.value_to_index(value), index);
404                assert_eq!(qi.index_to_key(index), qi.value_to_key(value));
405                for key in qi
406                    .index_to_key(index)
407                    .take(1)
408                    .chain(qi.index_to_key(index).rev().take(1))
409                {
410                    assert_eq!(qi.key_to_value(key), value);
411                    assert_eq!(qi.key_to_index(key), index);
412                }
413            }
414        }
415    }
416}
competitive/algorithm/quotient_index.rs

competitive/algorithm/
quotient_index.rs
