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 hilbert_curve_order(mut x: usize, mut y: usize, m: usize) -> usize {
58            let n = 1usize << m;
59            let mut ord = 0usize;
60            for k in (0..m).rev() {
61                let rx = x >> k & 1;
62                let ry = y >> k & 1;
63                ord += (1 << k * 2) * (3 * rx ^ ry);
64                if ry == 0 {
65                    if rx == 1 {
66                        x = n - x - 1;
67                        y = n - y - 1;
68                    }
69                    ::std::mem::swap(&mut x, &mut y);
70                }
71            }
72            ord
73        }
74        let lr: &[(usize, usize)] = $lr;
75        let q = lr.len();
76        let maxv = lr.iter().map(|&(l, r)| l.max(r)).max().unwrap_or_default();
77        let mut m = 0usize;
78        while maxv >= 1 << m {
79            m += 1;
80        }
81        let mut idx: Vec<usize> = (0..q).collect();
82        let ord: Vec<_> = lr
83            .iter()
84            .map(|&(l, r)| hilbert_curve_order(l, r, m))
85            .collect();
86        idx.sort_unstable_by_key(|&i| ord[i]);
87        let (mut $l, mut $r) = (0usize, 0usize);
88        for &$a in idx.iter() {
89            let (nl, nr) = lr[$a];
90            while $l > nl {
91                $l -= 1;
92                let $dl: usize = $l;
93                $decl;
94            }
95            while $r < nr {
96                let $ir: usize = $r;
97                $incr;
98                $r += 1;
99            }
100            while $l < nl {
101                let $il: usize = $l;
102                $incl;
103                $l += 1;
104            }
105            while $r > nr {
106                $r -= 1;
107                let $dr: usize = $r;
108                $decr;
109            }
110            $answer;
111        }
112    }};
113}
114
115#[cfg(test)]
116mod tests {
117    use crate::{rand, tools::NotEmptySegment as Nes, tools::Xorshift};
118
119    #[test]
120    fn test_mo_algorithm() {
121        let mut rng = Xorshift::default();
122        for _ in 0..50 {
123            rand!(rng, n: 1..50, q: 1..100, a: [1i64..1000; n], lr: [Nes(n); q]);
124            let mut ans = 0;
125            let mut acc = 0;
126            mo_algorithm!(
127                &lr,
128                (l, r),
129                |i| acc -= a[i],
130                |i| acc += a[i],
131                |i| acc += a[i],
132                |i| acc -= a[i],
133                |i| ans += acc
134            );
135            let mut exp = 0;
136            for (l, r) in lr {
137                exp += a[l..r].iter().sum::<i64>();
138            }
139            assert_eq!(ans, exp);
140        }
141    }
142}
competitive/algorithm/mo_algorithm.rs

competitive/algorithm/
mo_algorithm.rs
