#![expect(dead_code)]
use crate::notes::Interval;
#[derive(Debug, Clone)]
pub struct IntervalTree<T> {
pub(crate) intervals: Vec<Interval<T>>,
index_mapping: Vec<usize>,
subtree_max_end: Vec<T>,
}
impl<T> Default for IntervalTree<T>
where
T: Ord + Copy + Default,
{
fn default() -> Self {
Self::new()
}
}
impl<T> IntervalTree<T>
where
T: Ord + Copy + Default,
{
pub fn new() -> Self {
Self {
intervals: vec![],
index_mapping: vec![],
subtree_max_end: vec![],
}
}
pub fn with_intervals(intervals: Vec<Interval<T>>) -> Self {
let len = intervals.len();
let mut tree = Self {
intervals,
index_mapping: Vec::with_capacity(len),
subtree_max_end: Vec::with_capacity(len),
};
tree.rebuild(|_, _| {});
tree
}
pub fn rebuild(&mut self, mut swap: impl FnMut(usize, usize)) {
let mut ordering = calculate_reorder(
&self.intervals,
&mut self.index_mapping,
mapping::fill_inorder_veb,
);
let len = ordering.len();
for i in 0..len {
let mut current = i;
while ordering[current] != current {
let next = ordering[current];
if ordering[next] != next {
self.intervals.swap(current, next);
swap(current, next);
}
ordering[current] = current;
current = next;
}
}
fn assign_max_end<T: Ord + Copy>(tree: &mut IntervalTree<T>, index: usize) {
let left_child = index * 2 + 1;
let right_child = index * 2 + 2;
let len = tree.intervals.len();
let real_index = tree.index_mapping[index];
let self_max = tree.intervals[real_index].end();
if left_child >= len && right_child >= len {
tree.subtree_max_end[real_index] = tree.intervals[real_index].end();
} else if left_child < len && right_child < len {
assign_max_end(tree, left_child);
assign_max_end(tree, right_child);
tree.subtree_max_end[real_index] = self_max
.max(tree.subtree_max_end[tree.index_mapping[left_child]])
.max(tree.subtree_max_end[tree.index_mapping[right_child]]);
} else if left_child < len {
assign_max_end(tree, left_child);
tree.subtree_max_end[real_index] =
self_max.max(tree.subtree_max_end[tree.index_mapping[left_child]]);
} else {
assign_max_end(tree, right_child);
tree.subtree_max_end[real_index] =
self_max.max(tree.subtree_max_end[tree.index_mapping[right_child]]);
}
}
self.subtree_max_end
.resize(self.intervals.len(), T::default());
if !self.subtree_max_end.is_empty() {
assign_max_end(self, 0);
}
}
pub fn query_intersecting(&self, interval: Interval<T>) -> impl Iterator<Item = usize> {
iteration::IntersectingIntervalIterator::new(self, interval)
}
pub(crate) fn health_check(&self) {
#[cfg(debug_assertions)]
{
let len = self.intervals.len();
assert_eq!(len, self.index_mapping.len());
assert_eq!(len, self.subtree_max_end.len());
fn is_bst<T: Ord + Copy>(tree: &IntervalTree<T>, index: usize) {
if index >= tree.intervals.len() {
return;
}
let left_child = index * 2 + 1;
let right_child = index * 2 + 2;
let self_val = tree.intervals[tree.index_mapping[index]].start();
let self_max = tree.subtree_max_end[tree.index_mapping[index]];
if left_child < tree.intervals.len() {
let left_child_val = tree.intervals[tree.index_mapping[left_child]].start();
let left_child_max = tree.subtree_max_end[tree.index_mapping[left_child]];
assert!(left_child_val <= self_val);
assert!(self_max >= left_child_max);
is_bst(tree, left_child);
}
if right_child < tree.intervals.len() {
let right_child_val = tree.intervals[tree.index_mapping[right_child]].start();
let right_child_max = tree.subtree_max_end[tree.index_mapping[right_child]];
assert!(right_child_val >= self_val);
assert!(self_max >= right_child_max);
is_bst(tree, right_child);
}
}
is_bst(self, 0);
}
}
}
fn calculate_reorder<T: Ord + Copy>(
intervals: &[Interval<T>],
index_mapping: &mut Vec<usize>,
mapper: fn(&mut [usize]),
) -> Vec<usize> {
let mut sorted_indices = (0..intervals.len()).collect::<Vec<usize>>();
sorted_indices.sort_by_key(|i| intervals[*i].start());
let mut reorder = {
let mut mapping = calc_bfs_mapping(intervals.len(), mapper);
std::mem::swap(&mut mapping, index_mapping);
mapping
};
reorder.resize(intervals.len(), 0);
if reorder.is_empty() {
return reorder;
}
fn build_bst(
reorder: &mut [usize],
sorted_indices: &[usize],
map: &[usize],
node: usize,
counter: &mut usize,
) {
let left = node * 2 + 1;
let right = node * 2 + 2;
if left < reorder.len() {
build_bst(reorder, sorted_indices, map, left, counter);
}
reorder[map[node]] = sorted_indices[*counter];
*counter += 1;
if right < reorder.len() {
build_bst(reorder, sorted_indices, map, right, counter);
}
}
build_bst(&mut reorder, &sorted_indices, index_mapping, 0, &mut 0);
reorder
}
fn calc_bfs_mapping(n: usize, mapper: fn(&mut [usize])) -> Vec<usize> {
let mut map = vec![0; n];
mapper(&mut map);
map
}
pub mod iteration {
use super::*;
pub struct IntersectingIntervalIterator<'a, T> {
len: usize,
tree: &'a IntervalTree<T>,
querying_interval: Interval<T>,
current_bfs_node: usize,
terminate_at_right_of: usize,
}
impl<'a, T: Ord + Copy> IntersectingIntervalIterator<'a, T> {
pub(super) fn new(tree: &'a IntervalTree<T>, query: Interval<T>) -> Self {
let len = tree.intervals.len();
let mut v = Self {
len,
tree,
querying_interval: query,
current_bfs_node: 0,
terminate_at_right_of: 0,
};
if v.len > 0 {
let first_node = if tree.intervals[tree.index_mapping[0]].intersect(&query) {
0
} else if let Some(node) = v.next_node(0) {
node
} else {
v.len = 0;
0
};
v.current_bfs_node = first_node;
}
v
}
fn next_node(&mut self, node: usize) -> Option<usize> {
fn is_right_of(mut reference: usize, mut val: usize) -> bool {
if reference == 0 || reference == val {
return false;
}
let depth_ref = (usize::BITS - 1) - reference.leading_zeros();
let depth_val = (usize::BITS - 1) - val.leading_zeros();
if depth_ref > depth_val {
reference >>= depth_ref - depth_val;
} else if depth_val > depth_ref {
val >>= depth_val - depth_ref;
}
let diff_at = (usize::BITS - 1) - (reference ^ val).leading_zeros();
(val >> diff_at) & 1 == 1
}
let mut current_node = node + 1;
loop {
let left_child = current_node << 1;
let continue_at = if left_child <= self.len {
let real_index = self.tree.index_mapping[left_child - 1];
if self.tree.intervals[real_index].intersect(&self.querying_interval) {
return Some(left_child - 1);
} else if self.tree.subtree_max_end[real_index] < self.querying_interval.start()
{
left_child
} else {
current_node = left_child;
continue;
}
} else {
current_node
};
let real_index = self.tree.index_mapping[continue_at - 1];
if self.tree.intervals[real_index].start() > self.querying_interval.end() {
self.terminate_at_right_of = continue_at;
}
let is_left_sibling = continue_at % 2 == 0;
if is_left_sibling {
let right = continue_at + 1;
if right <= self.len {
let right_real_index = self.tree.index_mapping[right - 1];
if self.tree.intervals[right_real_index].intersect(&self.querying_interval)
{
return Some(right - 1);
} else {
current_node = right;
continue;
}
} else {
let umove = (right >> right.trailing_ones()) + 1;
if umove == 1
|| umove > self.len
|| is_right_of(self.terminate_at_right_of, umove)
{
return None;
} else {
let umove_real_index = self.tree.index_mapping[umove - 1];
if self.tree.intervals[umove_real_index]
.intersect(&self.querying_interval)
{
return Some(umove - 1);
} else {
current_node = umove;
continue;
}
}
}
} else {
let umove = (continue_at >> continue_at.trailing_ones()) + 1;
if umove == 1
|| umove > self.len
|| is_right_of(self.terminate_at_right_of, umove)
{
return None;
} else {
let umove_real_index = self.tree.index_mapping[umove - 1];
if self.tree.intervals[umove_real_index].intersect(&self.querying_interval)
{
return Some(umove - 1);
} else {
current_node = umove;
continue;
}
}
}
}
}
}
impl<T> Iterator for IntersectingIntervalIterator<'_, T>
where
T: Ord + Copy,
{
type Item = usize;
fn next(&mut self) -> Option<Self::Item> {
if self.len == 0 {
return None;
}
let current_node = self.current_bfs_node;
if let Some(next_node) = self.next_node(self.current_bfs_node) {
self.current_bfs_node = next_node;
} else {
self.len = 0;
}
Some(self.tree.index_mapping[current_node])
}
}
}
pub mod mapping {
use super::*;
pub fn fill_preorder_bfs(indices: &mut [usize]) {
indices.iter_mut().enumerate().for_each(|(i, e)| *e = i);
}
pub fn fill_preorder_veb(indices: &mut [usize]) {
fn inner(indices: &mut [usize], node: usize, height: usize, accumulate: &mut usize) {
let bottom_height = height.div_ceil(2);
let top_height = height - bottom_height;
if height == 1 {
indices[node] = *accumulate;
*accumulate += 1;
} else if height == 2 {
let left_child = node * 2 + 1;
let right_child = node * 2 + 2;
if node < indices.len() {
indices[node] = *accumulate;
*accumulate += 1;
}
if left_child < indices.len() {
indices[left_child] = *accumulate;
*accumulate += 1;
}
if right_child < indices.len() {
indices[right_child] = *accumulate;
*accumulate += 1;
}
} else {
let num_subtrees_same_depth = 1 << top_height;
let subtrees_begin_at = (node + 1) * (1 << top_height) - 1;
inner(indices, node, top_height, accumulate);
for i in 0..(num_subtrees_same_depth / 2) {
let subtree_root = subtrees_begin_at + i;
inner(indices, subtree_root, bottom_height, accumulate);
}
for i in (num_subtrees_same_depth / 2)..num_subtrees_same_depth {
let subtree_root = subtrees_begin_at + i;
inner(indices, subtree_root, bottom_height, accumulate);
}
}
}
let mut accumulate = 0;
if indices.is_empty() {
return;
}
if indices.len() == 1 {
indices[0] = 0;
return;
}
inner(indices, 0, fls(indices.len()), &mut accumulate);
}
pub fn fill_inorder_veb(indices: &mut [usize]) {
fn inner(indices: &mut [usize], node: usize, height: usize, accumulate: &mut usize) {
let bottom_height = height.div_ceil(2);
let top_height = height - bottom_height;
if height == 1 {
indices[node] = *accumulate;
*accumulate += 1;
} else if height == 2 {
let left_child = node * 2 + 1;
let right_child = node * 2 + 2;
if left_child < indices.len() {
indices[left_child] = *accumulate;
*accumulate += 1;
}
if node < indices.len() {
indices[node] = *accumulate;
*accumulate += 1;
}
if right_child < indices.len() {
indices[right_child] = *accumulate;
*accumulate += 1;
}
} else {
let num_subtrees_same_depth = 1 << top_height;
let subtrees_begin_at = (node + 1) * (1 << top_height) - 1;
for i in 0..(num_subtrees_same_depth / 2) {
let subtree_root = subtrees_begin_at + i;
inner(indices, subtree_root, bottom_height, accumulate);
}
inner(indices, node, top_height, accumulate);
for i in (num_subtrees_same_depth / 2)..num_subtrees_same_depth {
let subtree_root = subtrees_begin_at + i;
inner(indices, subtree_root, bottom_height, accumulate);
}
}
}
let mut accumulate = 0;
if indices.is_empty() {
return;
}
if indices.len() == 1 {
indices[0] = 0;
return;
}
inner(indices, 0, fls(indices.len()), &mut accumulate);
}
}
fn fls(mut f: usize) -> usize {
let mut order = 0;
while f != 0 {
f >>= 1;
order += 1;
}
order
}
fn hyperceil(f: usize) -> usize {
1 << fls(f - 1)
}
#[cfg(test)]
mod test {
use super::*;
use proptest::prelude::*;
#[test]
#[rustfmt::skip]
fn is_preorder_veb() {
let map =calc_bfs_mapping(63, mapping::fill_preorder_veb);
assert_eq!(map, vec![
1,
2, 5,
3, 4, 6, 7,
8, 15, 22, 29, 36, 43, 50, 57,
9, 12, 16, 19, 23, 26, 30, 33, 37, 40, 44, 47, 51, 54, 58, 61,
10, 11, 13, 14, 17, 18, 20, 21, 24, 25, 27, 28, 31, 32, 34, 35, 38, 39, 41, 42, 45, 46, 48, 49, 52, 53, 55, 56, 59, 60, 62, 63,
].into_iter().map(|x| x - 1).collect::<Vec<_>>())
}
#[test]
#[rustfmt::skip]
fn is_inorder_veb() {
let map = calc_bfs_mapping(63, mapping::fill_inorder_veb);
assert_eq!(map, vec![
32,
30, 34,
29, 31, 33, 35,
4, 11, 18, 25, 39, 46, 53, 60,
2, 6, 9, 13, 16, 20, 23, 27, 37, 41, 44, 48, 51, 55, 58, 62,
1, 3, 5, 7, 8, 10, 12, 14, 15, 17, 19, 21, 22, 24, 26, 28, 36, 38, 40, 42, 43, 45, 47, 49, 50, 52, 54, 56, 57, 59, 61, 63,
].into_iter().map(|x| x - 1).collect::<Vec<_>>())
}
proptest! {
#[test]
fn preorder_veb_mapping_is_0_n_permutation(n in 0usize..10000) {
let mut map = calc_bfs_mapping(n, mapping::fill_preorder_veb);
map.sort();
assert_eq!(map, (0..n).collect::<Vec<_>>());
}
}
proptest! {
#[test]
fn inorder_veb_mapping_is_0_n_permutation(n in 0usize..10000) {
let mut map = calc_bfs_mapping(n, mapping::fill_inorder_veb);
map.sort();
assert_eq!(map, (0..n).collect::<Vec<_>>());
}
}
#[test]
fn query_empty_tree() {
let tree = IntervalTree::<u32>::new();
let result = tree
.query_intersecting(Interval::new(0, 100))
.collect::<Vec<_>>();
assert!(result.is_empty())
}
fn arb_interval() -> impl Strategy<Value = Interval<u32>> {
(any::<u32>(), any::<u32>()).prop_map(|(v1, v2)| Interval::new(v1, v2))
}
fn arb_intervals(max_len: usize) -> impl Strategy<Value = Vec<Interval<u32>>> {
prop::collection::vec(arb_interval(), 0..max_len)
}
proptest! {
#[test]
fn reorder_should_be_0_n_permutation(intervals in arb_intervals(10000)) {
let mut map = vec![];
let mut reorder = calculate_reorder(&intervals, &mut map, mapping::fill_inorder_veb);
reorder.sort();
for (i, v) in reorder.iter().enumerate() {
assert_eq!(i, *v);
}
}
}
proptest! {
#[test]
fn can_build_arb_tree(intervals in arb_intervals(10000)) {
IntervalTree::with_intervals(intervals.clone()).health_check();
}
}
fn test_query_intervals(query: Interval<u32>, intervals: Vec<Interval<u32>>) {
let tree = IntervalTree::with_intervals(intervals.clone());
tree.health_check();
let tree_result = tree.query_intersecting(query).collect::<Vec<_>>();
let mut tree_items = tree_result
.into_iter()
.map(|i| tree.intervals[i])
.collect::<Vec<_>>();
let mut naive_items = intervals
.into_iter()
.filter(|iv| iv.intersect(&query))
.collect::<Vec<_>>();
tree_items.sort_by_key(|iv| (iv.start(), iv.end()));
naive_items.sort_by_key(|iv| (iv.start(), iv.end()));
assert_eq!(tree_items, naive_items);
}
proptest! {
#[test]
fn query_arb_intervals(query in arb_interval(), intervals in arb_intervals(100)) {
test_query_intervals(query, intervals);
}
}
proptest! {
#[ignore]
#[test]
fn query_arb_intervals_large(query in arb_interval(), intervals in arb_intervals(10000)) {
test_query_intervals(query, intervals);
}
}
#[test]
fn query_intervals_proptest_found_1() {
let intervals = vec![
Interval::new(0, 0),
Interval::new(0, 0),
Interval::new(1, 1),
];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_2() {
let intervals = vec![Interval::new(1, 1)];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_3() {
let intervals = vec![Interval::new(0, 0), Interval::new(1, 1)];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_4() {
let intervals = vec![Interval::new(1, 1), Interval::new(0, 0)];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_5() {
let intervals = vec![Interval::new(1, 1), Interval::new(0, 0)];
let query = Interval::new(1, 1);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_6() {
let intervals = vec![Interval::new(0, 0), Interval::new(1, 1)];
let query = Interval::new(1, 1);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_7() {
let intervals = vec![Interval::new(1, 1), Interval::new(1, 1)];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_8() {
let intervals = vec![
Interval::new(1, 1),
Interval::new(0, 0),
Interval::new(0, 0),
];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_9() {
let intervals = vec![
Interval::new(0, 0),
Interval::new(1, 1),
Interval::new(0, 0),
];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_10() {
let intervals = vec![
Interval::new(0, 0),
Interval::new(0, 0),
Interval::new(1, 1),
];
let query = Interval::new(0, 0);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_11() {
let intervals = vec![
Interval::new(0, 1),
Interval::new(0, 0),
Interval::new(0, 2),
Interval::new(0, 0),
Interval::new(0, 0),
];
let query = Interval::new(2, 3);
test_query_intervals(query, intervals);
}
#[test]
fn query_intervals_proptest_found_12() {
let intervals = vec![
Interval::new(0, 0),
Interval::new(0, 0),
Interval::new(0, 0),
Interval::new(1, 2),
Interval::new(1, 1),
Interval::new(0, 0),
];
let query = Interval::new(2, 2);
test_query_intervals(query, intervals);
}
#[test]
fn can_build_tree_all_0s() {
let intervals = vec![Interval::new(0, 0); 9];
IntervalTree::<u32>::with_intervals(intervals);
}
}