Mo’s Algorithm

Mo’s algorithm for offline range queries.

Mo’s algorithm solves queries of the form “given a value array and a list of range queries, answer each query using only O(|hi - lo|) moves between consecutive queries (with O(1) amortized per move) once the queries have been reordered by Mo’s block ordering.

The classic instance is counting distinct elements in a range. This crate ships count_distinct_in_range which handles that case directly. For more elaborate query types build your own solver using solve_blocked.

The Rust crate lives at mo_algorithm/src/lib.rs; build & run it with:

$ cd mo_algorithm && cargo run
$ cd mo_algorithm && cargo test

Generated API reference for the Mo’s algorithm crate. View the rendered API reference →

  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}