competitive/algorithm/
mo_algorithm.rs

1/// solve with Mo's algorithm
2///
3/// arg:
4/// - lr: slice of pair of usize
5/// - (l, r): ident of current pair
6/// - incl: |i| expr: del i, l+=1 (post)
7/// - decl: |i| expr: add i, l-=1 (pre)
8/// - incr: |i| expr: add i, r+=1 (post)
9/// - decr: |i| expr: del i, r-=1 (pre)
10/// - answer: |i| expr: answer i-th pair of lr
11///
12/// incr and decr can be omitted, if simultaneous
13///
14/// ```
15/// # use competitive::mo_algorithm;
16/// let (a, lr) = ([1, 2, 3], [(0, 1), (0, 2), (1, 3)]);
17/// let (mut ans, mut acc) = (0, 0);
18/// mo_algorithm!(
19///     &lr,
20///     (l, r),
21///     |i| acc -= a[i],
22///     |i| acc += a[i],
23///     |i| acc += a[i],
24///     |i| acc -= a[i],
25///     |i| ans += acc
26/// );
27/// assert_eq!(ans, 9);
28/// ```
29#[macro_export]
30macro_rules! mo_algorithm {
31    (
32        $lr:expr,
33        ($l:ident, $r:ident),
34        |$i:tt| $inc:expr,
35        |$d:tt| $dec:expr,
36        |$a:tt| $answer:expr $(,)?
37    ) => {{
38        $crate::mo_algorithm!(
39            $lr,
40            ($l, $r),
41            |$i| $inc,
42            |$d| $dec,
43            |$d| $dec,
44            |$i| $inc,
45            |$a| $answer
46        );
47    }};
48    (
49        $lr:expr,
50        ($l:ident, $r:ident),
51        |$il:tt| $incl:expr,
52        |$dl:tt| $decl:expr,
53        |$ir:tt| $incr:expr,
54        |$dr:tt| $decr:expr,
55        |$a:tt| $answer:expr $(,)?
56    ) => {{
57        fn mo_order<const SHIFTED: bool>(
58            lr: &[(usize, usize)],
59            maxv: usize,
60            width: usize,
61        ) -> (usize, Vec<usize>) {
62            let shift = usize::from(SHIFTED) * (width / 2);
63            let bucket = |x: usize| (x + shift) / width;
64            let buckets = bucket(maxv) + 1;
65            let mut pos = vec![0usize; buckets + 1];
66            for &(l, _) in lr {
67                pos[bucket(l) + 1] += 1;
68            }
69            for i in 1..=buckets {
70                pos[i] += pos[i - 1];
71            }
72            let mut idx = vec![0usize; lr.len()];
73            for (i, &(l, _)) in lr.iter().enumerate() {
74                let p = &mut pos[bucket(l)];
75                idx[*p] = i;
76                *p += 1;
77            }
78            idx[..pos[0]].sort_unstable_by_key(|&i| lr[i].1);
79            for b in (2..buckets).step_by(2) {
80                idx[pos[b - 1]..pos[b]].sort_unstable_by_key(|&i| lr[i].1);
81            }
82            for b in (1..buckets).step_by(2) {
83                idx[pos[b - 1]..pos[b]].sort_unstable_by_key(|&i| ::std::cmp::Reverse(lr[i].1));
84            }
85            let (mut l, mut r, mut len) = (0usize, 0usize, 0usize);
86            for &i in &idx {
87                let (nl, nr) = lr[i];
88                len += l.abs_diff(nl) + r.abs_diff(nr);
89                l = nl;
90                r = nr;
91            }
92            (len, idx)
93        }
94        let lr: &[(usize, usize)] = $lr;
95        let maxv = lr.iter().map(|&(l, r)| l.max(r)).max().unwrap_or_default();
96        let width = ((maxv as f64) / (lr.len().max(1) as f64).sqrt())
97            .round()
98            .max(1.0) as usize;
99        let (len0, idx0) = mo_order::<false>(lr, maxv, width);
100        let (len1, idx1) = mo_order::<true>(lr, maxv, width);
101        let idx = if len0 <= len1 { idx0 } else { idx1 };
102        let (mut $l, mut $r) = (0usize, 0usize);
103        for &$a in &idx {
104            let (nl, nr) = lr[$a];
105            while $l > nl {
106                $l -= 1;
107                let $dl: usize = $l;
108                $decl;
109            }
110            while $r < nr {
111                let $ir: usize = $r;
112                $incr;
113                $r += 1;
114            }
115            while $l < nl {
116                let $il: usize = $l;
117                $incl;
118                $l += 1;
119            }
120            while $r > nr {
121                $r -= 1;
122                let $dr: usize = $r;
123                $decr;
124            }
125            $answer;
126        }
127    }};
128}
129
130#[cfg(test)]
131mod tests {
132    use crate::{rand, tools::NotEmptySegment as Nes, tools::Xorshift};
133
134    #[test]
135    fn test_mo_algorithm() {
136        let mut rng = Xorshift::default();
137        for _ in 0..50 {
138            rand!(rng, n: 1..50, q: 1..100, a: [1i64..1000; n], lr: [Nes(n); q]);
139            let mut ans = 0;
140            let mut acc = 0;
141            mo_algorithm!(
142                &lr,
143                (l, r),
144                |i| acc -= a[i],
145                |i| acc += a[i],
146                |i| acc += a[i],
147                |i| acc -= a[i],
148                |i| ans += acc
149            );
150            let mut exp = 0;
151            for (l, r) in lr {
152                exp += a[l..r].iter().sum::<i64>();
153            }
154            assert_eq!(ans, exp);
155        }
156    }
157}
competitive/algorithm/mo_algorithm.rs

competitive/algorithm/
mo_algorithm.rs
