File size: 2,625 Bytes
2409829 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
use crate::aabb::Aabb;
use std::collections::HashSet;
pub struct QuadTree<T> {
bounding_box: Aabb,
depth: usize,
inner_node_capacity: usize,
subtrees: Option<Box<[QuadTree<T>; 4]>>,
pairs: Vec<(Aabb, T)>,
}
impl<T: Clone> QuadTree<T> {
pub fn new(bounding_box: Aabb, depth: usize, inner_node_capacity: usize) -> Self {
QuadTree {
bounding_box,
depth,
inner_node_capacity,
subtrees: None,
pairs: Vec::new(),
}
}
pub fn insert(&mut self, bounding_box: Aabb, value: T) -> bool {
if !crate::aabb::bounding_boxes_overlap(&bounding_box, &self.bounding_box) {
return false;
}
if self.depth > 0 && self.pairs.len() >= self.inner_node_capacity {
self.ensure_subtrees();
for tree in self.subtrees.as_mut().unwrap().iter_mut() {
tree.insert(bounding_box, value.clone());
}
} else {
self.pairs.push((bounding_box, value));
}
true
}
pub fn find(&self, bounding_box: &Aabb) -> HashSet<T>
where
T: Eq + std::hash::Hash + Clone,
{
let mut set = HashSet::new();
self.find_internal(bounding_box, &mut set);
set
}
fn find_internal(&self, bounding_box: &Aabb, set: &mut HashSet<T>)
where
T: Eq + std::hash::Hash + Clone,
{
if !crate::aabb::bounding_boxes_overlap(bounding_box, &self.bounding_box) {
return;
}
for (key, value) in &self.pairs {
if crate::aabb::bounding_boxes_overlap(bounding_box, key) {
set.insert(value.clone());
}
}
if let Some(subtrees) = &self.subtrees {
for tree in subtrees.iter() {
tree.find_internal(bounding_box, set);
}
}
}
fn ensure_subtrees(&mut self) {
if self.subtrees.is_some() {
return;
}
let mid_x = (self.bounding_box.left + self.bounding_box.right) / 2.;
let mid_y = (self.bounding_box.top + self.bounding_box.bottom) / 2.;
self.subtrees = Some(Box::new([
QuadTree::new(
Aabb {
top: self.bounding_box.top,
right: mid_x,
bottom: mid_y,
left: self.bounding_box.left,
},
self.depth - 1,
self.inner_node_capacity,
),
QuadTree::new(
Aabb {
top: self.bounding_box.top,
right: self.bounding_box.right,
bottom: mid_y,
left: mid_x,
},
self.depth - 1,
self.inner_node_capacity,
),
QuadTree::new(
Aabb {
top: mid_y,
right: mid_x,
bottom: self.bounding_box.bottom,
left: self.bounding_box.left,
},
self.depth - 1,
self.inner_node_capacity,
),
QuadTree::new(
Aabb {
top: mid_y,
right: self.bounding_box.right,
bottom: self.bounding_box.bottom,
left: mid_x,
},
self.depth - 1,
self.inner_node_capacity,
),
]));
}
}
|