Skip to main content

competitive/data_structure/
range_minimum_query.rs

1const BLOCK_SIZE: usize = 64;
2const BLOCK_SHIFT: usize = BLOCK_SIZE.trailing_zeros() as usize;
3const BLOCK_MASK: usize = BLOCK_SIZE - 1;
4
5#[derive(Clone, Debug)]
6pub struct RangeMinimumQuery<T> {
7    data: Vec<T>,
8    mask: Vec<u64>,
9    sparse: Vec<usize>,
10    sparse_offset: Vec<usize>,
11}
12
13impl<T> RangeMinimumQuery<T>
14where
15    T: Ord,
16{
17    pub fn new(data: Vec<T>) -> Self {
18        let n = data.len();
19        let blocks = (n + BLOCK_MASK) >> BLOCK_SHIFT;
20        let mut mask = vec![0u64; n];
21        let mut block_min = vec![0usize; blocks];
22
23        let mut stack: Vec<usize> = Vec::with_capacity(BLOCK_SIZE);
24        for (block, min_slot) in block_min.iter_mut().enumerate() {
25            stack.clear();
26            let start = block << BLOCK_SHIFT;
27            let end = (start + BLOCK_SIZE).min(n);
28            for i in start..end {
29                while let Some(&top) = stack.last() {
30                    if data[top] < data[i] {
31                        break;
32                    }
33                    stack.pop();
34                }
35                let prev_mask = stack.last().map_or(0, |&top| mask[top]);
36                mask[i] = prev_mask | 1u64 << (i - start);
37                stack.push(i);
38            }
39            *min_slot = (start..end).min_by_key(|&i| &data[i]).unwrap();
40        }
41
42        let mut total = 0usize;
43        let mut k = 0usize;
44        while (1usize << k) <= blocks {
45            total += blocks - (1usize << k) + 1;
46            k += 1;
47        }
48        let mut sparse = Vec::with_capacity(total);
49        let mut sparse_offset = Vec::with_capacity(k);
50        sparse_offset.push(0);
51        sparse.extend_from_slice(&block_min);
52        let mut k = 1usize;
53        let mut prev_offset = 0usize;
54        while (1usize << k) <= blocks {
55            let len = blocks - (1usize << k) + 1;
56            let offset = sparse.len();
57            sparse_offset.push(offset);
58            let jump = 1usize << (k - 1);
59            for i in 0..len {
60                let left = sparse[prev_offset + i];
61                let right = sparse[prev_offset + i + jump];
62                sparse.push(if data[left] < data[right] {
63                    left
64                } else {
65                    right
66                });
67            }
68            prev_offset = offset;
69            k += 1;
70        }
71
72        Self {
73            data,
74            mask,
75            sparse,
76            sparse_offset,
77        }
78    }
79
80    pub fn len(&self) -> usize {
81        self.data.len()
82    }
83
84    pub fn is_empty(&self) -> bool {
85        self.data.is_empty()
86    }
87
88    pub fn get(&self, index: usize) -> &T {
89        &self.data[index]
90    }
91
92    fn min_index(&self, a: usize, b: usize) -> usize {
93        if self.data[a] < self.data[b] { a } else { b }
94    }
95
96    fn argmin_in_block(&self, l: usize, r: usize) -> usize {
97        debug_assert!(l <= r);
98        debug_assert!(l >> BLOCK_SHIFT == r >> BLOCK_SHIFT);
99        let start = l & !BLOCK_MASK;
100        let shift = l & BLOCK_MASK;
101        let masked = self.mask[r] & (!0u64 << shift);
102        let offset = masked.trailing_zeros() as usize;
103        start + offset
104    }
105
106    fn argmin_blocks(&self, l_block: usize, r_block: usize) -> usize {
107        debug_assert!(l_block < r_block);
108        let len = r_block - l_block;
109        let k = len.ilog2() as usize;
110        let offset = self.sparse_offset[k];
111        let left = self.sparse[offset + l_block];
112        let right = self.sparse[offset + r_block - (1usize << k)];
113        self.min_index(left, right)
114    }
115
116    pub fn argmin(&self, l: usize, r: usize) -> usize {
117        assert!(l < r);
118        assert!(r <= self.data.len());
119        let r = r - 1;
120        let bl = l >> BLOCK_SHIFT;
121        let br = r >> BLOCK_SHIFT;
122        if bl == br {
123            return self.argmin_in_block(l, r);
124        }
125        let left_end = ((bl + 1) << BLOCK_SHIFT) - 1;
126        let right_start = br << BLOCK_SHIFT;
127        let mut best = self.min_index(
128            self.argmin_in_block(l, left_end),
129            self.argmin_in_block(right_start, r),
130        );
131        if bl + 1 < br {
132            let mid = self.argmin_blocks(bl + 1, br);
133            best = self.min_index(best, mid);
134        }
135        best
136    }
137}
138
139impl<T> RangeMinimumQuery<T>
140where
141    T: Ord + Clone,
142{
143    pub fn fold(&self, l: usize, r: usize) -> T {
144        let idx = self.argmin(l, r);
145        self.data[idx].clone()
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::{
153        rand,
154        tools::{NotEmptySegment as Nes, Xorshift},
155    };
156
157    #[test]
158    fn test_range_minimum_query() {
159        let mut rng = Xorshift::default();
160        for _ in 0..100 {
161            rand!(rng, n: 1..200, arr: [-1000i64..=1000; n]);
162            let rmq = RangeMinimumQuery::new(arr.clone());
163            for _ in 0..200 {
164                rand!(rng, (l, r): Nes(n));
165                let expected = arr[l..r].iter().min().cloned().unwrap();
166                assert_eq!(rmq.fold(l, r), expected);
167                let idx = rmq.argmin(l, r);
168                assert!(l <= idx && idx < r);
169                assert_eq!(rmq.get(idx), &expected);
170            }
171        }
172    }
173}