Range in a BTreeMap with tuple keys

Given a BTreeMap with tuple keys, I want to make queries like "find all tuples with first value = x" or "find all with last value = y". However, for the last value, the range isn't exactly working. range((1, 3)..=(3, 3)) also finds (2, 1) because it falls between (1, 3) and (2, 3).

Link to Playground.

Tuples are always ordered lexicographically, so you can't have the first element as the fastest-changing index. Your query isn't expressible as-is.

You can build a list of all possible keys instead:

    for key in (1..=3).map(|x| (x, 3)) {
        if let Some(value) = map.get(&key) {
            println!("{:?}: {:?}", key, value);
        }
    }

BTreeMap doesn't support 2-dimensional range queries. What you're looking for is a 2D range tree. I don't know if there is an implementation available already.

3 Likes

It may be less efficient than using a data structure designed for this, but is this what you want?

1 Like

If you want to filter on only one of the tuple fields at a time, you can (I guess) pass .. for the field on which you aren't filtering or have separate map_range_... functions.

Alternatively, if you don't like the O(n) lookup for the case in question, you can have two BTreeMaps: one with the keys as-is, and one with their elements reversed.

The implementation you provided is a good example for a Rust freshman like me, but I came up with the following data structure that serves my purpose. I'm, of course, open to any suggestions as to how to improve it.

use std::collections::{BTreeSet, HashMap};

pub struct LRIndex<'a> {
    idx: HashMap<u8, BTreeSet<usize>>,
    dominoes: &'a [(u8, u8)],
}

#[derive(Debug, Clone, PartialEq)]
pub enum MatchSide {
    L,
    R,
}

// Lifetime in structs: https://stackoverflow.com/a/27590535/839733
impl<'a> LRIndex<'a> {
    pub fn new(dominoes: &'a [(u8, u8)]) -> Self {
        let mut idx: HashMap<u8, BTreeSet<usize>> = HashMap::new();

        for (i, (l, r)) in dominoes.iter().enumerate() {
            idx.entry(*l).or_default().insert(i);
            idx.entry(*r).or_default().insert(i);
        }
        Self { idx, dominoes }
    }

    fn is_used(&self, i: usize) -> bool {
        assert!(
            (0..self.dominoes.len()).contains(&i),
            "Index out of bounds: {i}"
        );
        !self.idx[&self.dominoes[i].0].contains(&i)
    }

    pub fn acquire(&mut self, i: usize) {
        assert!(!self.is_used(i), "{i} is already in use");
        self.a_r(i, |e| {
            e.remove(&i);
        });
    }

    pub fn release(&mut self, i: usize) {
        assert!(self.is_used(i), "{i} is not used");
        self.a_r(i, |e| {
            e.insert(i);
        });
    }

    fn a_r<F>(&mut self, i: usize, f: F)
    where
        F: FnOnce(&mut BTreeSet<usize>) + Copy,
    {
        let (l, r) = self.dominoes[i];
        self.idx.entry(l).and_modify(|e| {
            f(e);
        });
        self.idx.entry(r).and_modify(|e| {
            f(e);
        });
    }

    // 1. Should return all indices that have either left
    //    or right matching the the given domino side.
    // 2. Should not contain the given dominoes[i].
    // 3. Should not contain duplicate indices.
    // 4. Should contain identical dominoes with unique indices.
    pub fn candidates(&self, i: usize, side: MatchSide) -> ((u8, u8), Vec<(usize, MatchSide)>) {
        assert!(
            (0..self.dominoes.len()).contains(&i),
            "Index out of bounds: {i}"
        );
        let (left, right) = self.dominoes[i];
        let k = if side == MatchSide::R { right } else { left };
        let c = self.idx[&k]
            .iter()
            .filter(|x| **x != i)
            .map(|x| {
                let other = self.dominoes[*x];
                let side = if k == other.0 {
                    MatchSide::L
                } else {
                    MatchSide::R
                };
                (*x, side)
            })
            .collect();
        ((left, right), c)
    }
}