Interval Tree

extents.txt

-1 20
3 5
8 9

interval_tree.rs

use std::cell::RefCell;
use std::fmt::Debug;
use std::ops::{Add, Div};
use std::rc::Rc;

type NodePtr<T> = Rc<RefCell<Node<T>>>;

#[derive(Debug)]
pub struct IntervalTree<T: Debug> {
    root: NodePtr<T>,
}

#[derive(Debug)]
struct Node<T: Debug> {
    center: T,
    left: Option<NodePtr<T>>,
    right: Option<NodePtr<T>>,
    overlap_by_start: Vec<(T, T)>,
    overlap_by_end: Vec<(T, T)>,
}

impl<T: Debug> Into<NodePtr<T>> for Node<T> {
    fn into(self) -> NodePtr<T> {
        Rc::new(RefCell::new(self))
    }
}

impl<T> Node<T>
where
    T: Copy + Debug + Add<Output = T> + Div<Output = T> + Ord + From<i64> + std::iter::Sum,
{
    fn new(intervals: &[(T, T)]) -> Node<T> {
        let count: T = T::from(intervals.len() as i64);
        let center = intervals
            .iter()
            .map(|&(s, e)| (s + e) / T::from(2))
            .sum::<T>()
            / count;
        let (left, overlap_right): (Vec<(T, T)>, Vec<(T, T)>) = intervals
            .into_iter()
            .partition(|(s, e)| *s < center && *e < center);
        let (right, overlap): (Vec<(T, T)>, Vec<(T, T)>) = overlap_right
            .iter()
            .partition(|(s, e)| *s > center && *e > center);
        let mut overlap_by_end = overlap.clone();
        let mut overlap_by_start = overlap;
        overlap_by_start.sort_by_key(|&(s, _)| s);
        overlap_by_end.sort_by_key(|&(_, e)| e);
        overlap_by_end.reverse();

        Node {
            center,
            left: if left.len() > 0 {
                Some(Node::new(&left).into())
            } else {
                None
            },
            right: if right.len() > 0 {
                Some(Node::new(&right).into())
            } else {
                None
            },
            overlap_by_start,
            overlap_by_end,
        }
    }

    fn intersect_point(&self, p: T, intervals: &mut Vec<(T, T)>) {
        if p < self.center {
            intervals.extend(self.overlap_by_start.iter().take_while(|(s, _)| *s <= p));
            if let Some(ref left) = self.left {
                left.borrow().intersect_point(p, intervals);
            }
        } else if p > self.center {
            intervals.extend(self.overlap_by_end.iter().take_while(|(_, e)| *e >= p));
            if let Some(ref right) = self.right {
                right.borrow().intersect_point(p, intervals);
            }
        } else {
            intervals.extend_from_slice(&self.overlap_by_start);
        }
    }
}

impl<T> IntervalTree<T>
where
    T: Copy + Debug + Add<Output = T> + Div<Output = T> + Ord + From<i64> + std::iter::Sum,
{
    pub fn new(intervals: &[(T, T)]) -> IntervalTree<T> {
        IntervalTree {
            root: Node::new(intervals).into(),
        }
    }

    pub fn intersect_point(&self, p: T) -> Vec<(T, T)> {
        let mut intervals = Vec::new();

        self.root.borrow().intersect_point(p, &mut intervals);

        intervals
    }
}

main.rs

mod interval_tree;

use crate::interval_tree::IntervalTree;
use std::io::prelude::*;

type IntType = i64;

fn main() -> std::io::Result<()> {
    let tree = {
        let intervals: Vec<(IntType, IntType)> = read_file_lines("extents.txt")?
            .map(|s| parse_extent(&s))
            .collect();
        IntervalTree::new(&intervals)
    };

    let points = read_file_lines("points.txt")?.map(|s| parse_point(&s));
    let num_intersections = points.map(|p| tree.intersect_point(p).len());

    for i in num_intersections {
        println!("{}", i);
    }

    Ok(())
}

fn read_file_lines(filename: &str) -> std::io::Result<impl Iterator<Item = String>> {
    std::fs::OpenOptions::new()
        .read(true)
        .open(filename)
        .map(std::io::BufReader::new)
        .map(|r| r.lines().map(|l| l.unwrap()))
}

fn parse_extent(s: &str) -> (IntType, IntType) {
    let mut i = s.split(" ");
    (
        i.next().unwrap().parse().unwrap(),
        i.next().unwrap().parse().unwrap(),
    )
}

fn parse_point(s: &str) -> IntType {
    s.parse().unwrap()
}

points.txt

25
4
7
9