mo_algorithm/
lib.rs

1//! Mo's algorithm for offline range queries.
2//!
3//! Mo's algorithm solves queries of the form "given a value array `a`, and
4//! a list of range queries `[lo, hi)`, answer each query using only
5//! `O(|hi - lo|)` moves between consecutive queries (with `O(1)` amortized
6//! per move) once the queries have been reordered.
7//!
8//! The classic instance is counting distinct elements in a range. This
9//! crate ships [`count_distinct_in_range`], which handles that case
10//! directly:
11//!
12//! ```no_run
13//! use mo_algorithm::{count_distinct_in_range, Query};
14//!
15//! let a = vec![1, 1, 2, 1, 3];
16//! let queries = vec![Query { lo: 0, hi: 3 }, Query { lo: 2, hi: 5 }];
17//! let answers = count_distinct_in_range(&a, &queries);
18//! assert_eq!(answers, vec![2, 3]);
19//! ```
20//!
21//! For more elaborate query types build your own solver using
22//! [`solve_blocked`].
23
24/// A single query on a half-open interval `[lo, hi)`.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct Query {
27    /// Inclusive lower bound of the query range.
28    pub lo: usize,
29    /// Exclusive upper bound of the query range.
30    pub hi: usize,
31}
32
33impl Query {
34    fn block(&self, block_size: usize) -> usize {
35        self.lo / block_size
36    }
37}
38
39/// Run a generic Mo's algorithm over `n` elements answering `queries`.
40///
41/// All moving window bookkeeping is done in this function; the caller
42/// supplies a single `state` value passed to each callback. Interior
43/// mutability (e.g. `RefCell`) lets one state object serve all three
44/// callbacks if you need shared mutable bookkeeping.
45pub fn solve_blocked<State, FAdd, FRemove, FVisit>(
46    n: usize,
47    queries: &[Query],
48    state: &mut State,
49    mut add: FAdd,
50    mut remove: FRemove,
51    mut visit: FVisit,
52) where
53    State: ?Sized,
54    FAdd: FnMut(&mut State, usize),
55    FRemove: FnMut(&mut State, usize),
56    FVisit: FnMut(&mut State, usize),
57{
58    let _ = n;
59    if queries.is_empty() {
60        return;
61    }
62    let block_size = (queries.len() as f64).sqrt().ceil() as usize;
63    assert!(block_size >= 1, "block size must be at least 1");
64
65    // Sort queries by Mo's ordering.
66    let mut order: Vec<usize> = (0..queries.len()).collect();
67    order.sort_by_key(|&i| {
68        let q = queries[i];
69        let b = q.block(block_size);
70        (b, if b & 1 == 0 { q.hi as i64 } else { -(q.hi as i64) })
71    });
72
73    let mut cur_lo = queries[order[0]].lo;
74    let mut cur_hi = queries[order[0]].lo;
75
76    for &qi in &order {
77        let q = queries[qi];
78        // Expand / contract window to cover [q.lo, q.hi).
79        while cur_lo > q.lo {
80            cur_lo -= 1;
81            add(state, cur_lo);
82        }
83        while cur_hi < q.hi {
84            add(state, cur_hi);
85            cur_hi += 1;
86        }
87        while cur_lo < q.lo {
88            remove(state, cur_lo);
89            cur_lo += 1;
90        }
91        while cur_hi > q.hi {
92            cur_hi -= 1;
93            remove(state, cur_hi);
94        }
95        visit(state, qi);
96    }
97}
98
99/// Solve "number of distinct values in range" queries using Mo's
100/// algorithm.
101///
102/// `a` is the value array; `queries` are the half-open intervals.
103/// Returns one answer per query.
104pub fn count_distinct_in_range<T: Eq + std::hash::Hash + Copy>(
105    a: &[T],
106    queries: &[Query],
107) -> Vec<usize> {
108    struct State<T> {
109        freq: std::collections::HashMap<T, usize>,
110        distinct: usize,
111    }
112
113    let mut state: State<T> = State {
114        freq: std::collections::HashMap::new(),
115        distinct: 0,
116    };
117    let mut answers: Vec<usize> = vec![0; queries.len()];
118
119    solve_blocked(
120        a.len(),
121        queries,
122        &mut state,
123        |st, idx| {
124            let entry = st.freq.entry(a[idx]).or_insert(0);
125            if *entry == 0 {
126                st.distinct += 1;
127            }
128            *entry += 1;
129        },
130        |st, idx| {
131            if let Some(v) = st.freq.get_mut(&a[idx]) {
132                *v -= 1;
133                if *v == 0 {
134                    st.distinct -= 1;
135                }
136            }
137        },
138        |st, qi| {
139            answers[qi] = st.distinct;
140        },
141    );
142
143    answers
144}
mo_algorithm/lib.rs

mo_algorithm/
lib.rs
