Range in a BTreeMap with tuple keys

[bruce-wayne]

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.

[H2CO3]

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);
        }
    }

[tczajka]

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.

[8573]

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

Rust Playground

[8573]

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.

[H2CO3]

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.

[bruce-wayne]

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)
    }
}