nekolib/algo/
mo.rs

1//! Mo's algorithm。
2
3use super::super::math::sqrt;
4use super::super::traits::elastic_slice;
5
6use std::ops::Range;
7
8use elastic_slice::{
9    ElasticSlice, ExpandBack, ExpandFront, ShrinkBack, ShrinkFront, SliceHash,
10};
11use sqrt::Sqrt;
12
13/// Mo's algorithm。
14///
15/// オフラインのクエリを $q$ 個処理する。
16/// $i$ 番目のクエリは「区間 $[l\_i, r\_i)$ と値 $x\_i$ によって計算される値を求めよ」
17/// ということを意味する。
18///
19/// # Complexity
20/// 区間の全長 $n$、クエリ数 $q$、パラメータ $b$ に対して、
21/// `shrink_front` と `expand_front` を高々 $q\\cdot b$ 回、
22/// `shrink_back` と `expand_back` を高々 $n\^2/b$ 回行う。
23///
24/// # Hints
25/// 相加・相乗平均の等号成立条件から $b = n\\cdot q\^{-1/2}$
26/// とするのがよさげに思えるが、実際には手でパラメータを調整したくなることも多い。
27/// そのため、引数に `None` を渡した場合は上記の値を用い、`Some(b)`
28/// を渡した場合はその `b` を用いる実装にしている。
29///
30/// 定数での除算は最適化により乗算などに置き換えられることが期待されるので、
31/// 予め最大ケースにおける $b$ を計算するなどして、その値を渡す方がいいかも。
32/// 個人的には $224$ から $384$ くらいの大きさの $32$ の倍数に祈ることが多い。
33///
34/// # Examples
35/// ```
36/// use std::collections::BTreeMap;
37///
38/// use nekolib::traits::{
39///     ElasticSlice, ExpandBack, ExpandFront, ShrinkBack,
40///     ShrinkFront, SliceHash,
41/// };
42/// use nekolib::algo::mo;
43///
44/// struct RangeDistinct {
45///     buf: Vec<i32>,
46///     start: usize,
47///     end: usize,
48///     count: BTreeMap<i32, usize>,
49/// }
50///
51/// impl From<Vec<i32>> for RangeDistinct {
52///     fn from(buf: Vec<i32>) -> Self {
53///         Self { buf, start: 0, end: 0, count: BTreeMap::new() }
54///     }
55/// }
56///
57/// impl ElasticSlice for RangeDistinct {
58///     fn reset(&mut self) {
59///         self.start = 0;
60///         self.end = 0;
61///         self.count.clear();
62///     }
63///     fn full_len(&self) -> usize { self.buf.len() }
64///     fn start(&self) -> usize { self.start }
65///     fn end(&self) -> usize { self.end }
66/// }
67///
68/// /// 区間 `start..end` に含まれる整数の集合と、`x` のみからなる
69/// /// 集合との和集合の要素数を返す。
70/// impl SliceHash for RangeDistinct {
71///     type Salt = i32;
72///     type Hashed = usize;
73///     fn hash(&self, x: i32) -> usize {
74///         self.count.len()
75///             + if self.count.contains_key(&x) { 0 } else { 1 }
76///     }
77/// }
78///
79/// impl ExpandBack for RangeDistinct {
80///     fn expand_back(&mut self) {
81///         let k = self.buf[self.end];
82///         *self.count.entry(k).or_insert(0) += 1;
83///         self.end += 1;
84///     }
85/// }
86///
87/// impl ExpandFront for RangeDistinct {
88///     fn expand_front(&mut self) {
89///         self.start -= 1;
90///         let k = self.buf[self.start];
91///         *self.count.entry(k).or_insert(0) += 1;
92///     }
93/// }
94///
95/// impl ShrinkBack for RangeDistinct {
96///     fn shrink_back(&mut self) {
97///         self.end -= 1;
98///         let k = self.buf[self.end];
99///         match self.count.get_mut(&k) {
100///             Some(x) if x == &1 => { self.count.remove(&k); }
101///             Some(x) => *x -= 1,
102///             None => unreachable!(),
103///         }
104///     }
105/// }
106///
107/// impl ShrinkFront for RangeDistinct {
108///     fn shrink_front(&mut self) {
109///         let k = self.buf[self.start];
110///         match self.count.get_mut(&k) {
111///             Some(x) if x == &1 => { self.count.remove(&k); }
112///             Some(x) => *x -= 1,
113///             None => unreachable!(),
114///         }
115///         self.start += 1;
116///     }
117/// }
118///
119/// let rd: RangeDistinct = vec![1, 4, 1, 4, 2, 1, 3, 5, 6].into();
120/// let qs = vec![(0..4, 1), (0..4, 2), (2..6, 1), (3..9, 2)];
121/// assert_eq!(mo(rd, qs, Some(4)), vec![2, 3, 3, 6]);
122/// ```
123pub fn mo<S>(
124    mut slice: S,
125    q: Vec<(Range<usize>, S::Salt)>,
126    b: Option<usize>,
127) -> Vec<S::Hashed>
128where
129    S: ElasticSlice
130        + ExpandFront
131        + ExpandBack
132        + ShrinkFront
133        + ShrinkBack
134        + SliceHash,
135    S::Hashed: Clone,
136{
137    let b = match b {
138        Some(b) => b,
139        None => 1.max(slice.len().sqrt()),
140    };
141    let qn = q.len();
142    let mut q: Vec<_> = q.into_iter().enumerate().collect();
143    q.sort_unstable_by_key(|(_, (ir, _))| {
144        let &Range { start: is, end: ie } = ir;
145        let ib = is / b;
146        if ib % 2 == 0 { (ib, ie) } else { (ib, !ie) }
147    });
148
149    let mut res = vec![None; qn];
150    slice.reset();
151    for (i, (Range { start: ql, end: qr }, x)) in q {
152        while slice.end() < qr {
153            slice.expand_back();
154        }
155        while slice.start() > ql {
156            slice.expand_front();
157        }
158        while slice.start() < ql {
159            slice.shrink_front();
160        }
161        while slice.end() > qr {
162            slice.shrink_back();
163        }
164        res[i] = Some(slice.hash(x));
165    }
166    res.into_iter().map(std::option::Option::unwrap).collect()
167}
nekolib/algo/mo.rs

nekolib/algo/
mo.rs
