haar_lib/algo/
mo.rs

1//! Mo's algorithm
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/dwacon2017-honsen/tasks/dwango2017final_b>
5
6use std::cmp::Ordering;
7
8/// Mo's algorithm
9pub struct Mo<'a> {
10    append_left: Box<dyn 'a + FnMut(usize)>,
11    append_right: Box<dyn 'a + FnMut(usize)>,
12    remove_left: Box<dyn 'a + FnMut(usize)>,
13    remove_right: Box<dyn 'a + FnMut(usize)>,
14    query: Box<dyn 'a + FnMut(usize)>,
15    q: usize,
16    width: usize,
17    left: Vec<usize>,
18    right: Vec<usize>,
19    ord: Vec<usize>,
20}
21
22impl<'a> Mo<'a> {
23    /// クエリ処理用の関数を登録して[`Mo`]を作る。
24    pub fn new(
25        n: usize,
26        append_left: Box<impl 'a + FnMut(usize)>,
27        append_right: Box<impl 'a + FnMut(usize)>,
28        remove_left: Box<impl 'a + FnMut(usize)>,
29        remove_right: Box<impl 'a + FnMut(usize)>,
30        query: Box<impl 'a + FnMut(usize)>,
31    ) -> Self {
32        Self {
33            append_left,
34            append_right,
35            remove_left,
36            remove_right,
37            query,
38            q: 0,
39            width: (n as f64).sqrt() as usize,
40            left: vec![],
41            right: vec![],
42            ord: vec![],
43        }
44    }
45
46    /// 範囲`[l, r)`を登録する。
47    pub fn add(&mut self, l: usize, r: usize) {
48        self.left.push(l);
49        self.right.push(r);
50        self.ord.push(self.q);
51        self.q += 1;
52    }
53
54    /// 登録した範囲をすべて訪れるように実行する。
55    pub fn run(mut self) {
56        let left = self.left;
57        let right = self.right;
58        let width = self.width;
59
60        self.ord.sort_by(|&i, &j| {
61            let a = left[i] / width;
62            let b = left[j] / width;
63
64            if a == b {
65                if a % 2 == 1 {
66                    right[i].cmp(&right[j])
67                } else {
68                    right[j].cmp(&right[i])
69                }
70            } else {
71                a.cmp(&b)
72            }
73        });
74
75        let mut l = left[self.ord[0]];
76        let mut r = left[self.ord[0]];
77
78        for id in self.ord {
79            let left = left[id];
80            let right = right[id];
81
82            while l != left || r != right {
83                match l.cmp(&left) {
84                    Ordering::Greater => {
85                        l -= 1;
86                        (self.append_left)(l);
87                    }
88                    Ordering::Less => {
89                        (self.remove_left)(l);
90                        l += 1;
91                    }
92                    _ => {}
93                }
94
95                match r.cmp(&right) {
96                    Ordering::Less => {
97                        (self.append_right)(r);
98                        r += 1;
99                    }
100                    Ordering::Greater => {
101                        r -= 1;
102                        (self.remove_right)(r);
103                    }
104                    _ => {}
105                }
106            }
107
108            (self.query)(id);
109        }
110    }
111}