use crate::query::inner_prelude::*;
use core::cmp::Ordering;
pub trait Knearest {
type T: Aabb<Num = Self::N>;
type N: Num;
fn distance_to_aaline<A: Axis>(
&mut self,
point: Vec2<Self::N>,
axis: A,
val: Self::N,
) -> Self::N;
fn distance_to_broad(&mut self, point: Vec2<Self::N>, a: PMut<Self::T>) -> Option<Self::N>;
fn distance_to_fine(&mut self, point: Vec2<Self::N>, a: PMut<Self::T>) -> Self::N;
}
pub fn default_rect_knearest<T: Aabb>(tree: &Tree<T>) -> impl Knearest<T = T, N = T::Num>
where
T::Num: num_traits::Signed + num_traits::Zero,
{
use num_traits::Signed;
use num_traits::Zero;
from_closure(
tree,
(),
|_, _, _| None,
|_, point, a| {
a.get()
.distance_squared_to_point(point)
.unwrap_or_else(T::Num::zero)
},
|_, point, a| (point.x - a).abs() * (point.x - a).abs(),
|_, point, a| (point.y - a).abs() * (point.y - a).abs(),
)
}
struct KnearestBorrow<'a, K>(&'a mut K);
impl<'a, K: Knearest> Knearest for KnearestBorrow<'a, K> {
type T = K::T;
type N = K::N;
fn distance_to_aaline<A: Axis>(
&mut self,
point: Vec2<Self::N>,
axis: A,
val: Self::N,
) -> Self::N {
self.0.distance_to_aaline(point, axis, val)
}
fn distance_to_broad(&mut self, point: Vec2<Self::N>, rect: PMut<Self::T>) -> Option<Self::N> {
self.0.distance_to_broad(point, rect)
}
fn distance_to_fine(&mut self, point: Vec2<Self::N>, bot: PMut<Self::T>) -> Self::N {
self.0.distance_to_fine(point, bot)
}
}
use crate::Tree;
pub fn from_closure<Acc, T: Aabb>(
_tree: &Tree<T>,
acc: Acc,
broad: impl FnMut(&mut Acc, Vec2<T::Num>, PMut<T>) -> Option<T::Num>,
fine: impl FnMut(&mut Acc, Vec2<T::Num>, PMut<T>) -> T::Num,
xline: impl FnMut(&mut Acc, Vec2<T::Num>, T::Num) -> T::Num,
yline: impl FnMut(&mut Acc, Vec2<T::Num>, T::Num) -> T::Num,
) -> impl Knearest<T = T, N = T::Num> {
struct KnearestClosure<T: Aabb, Acc, B, C, D, E> {
pub _p: PhantomData<T>,
pub acc: Acc,
pub broad: B,
pub fine: C,
pub xline: D,
pub yline: E,
}
impl<'a, T: Aabb, Acc, B, C, D, E> Knearest for KnearestClosure<T, Acc, B, C, D, E>
where
B: FnMut(&mut Acc, Vec2<T::Num>, PMut<T>) -> Option<T::Num>,
C: FnMut(&mut Acc, Vec2<T::Num>, PMut<T>) -> T::Num,
D: FnMut(&mut Acc, Vec2<T::Num>, T::Num) -> T::Num,
E: FnMut(&mut Acc, Vec2<T::Num>, T::Num) -> T::Num,
{
type T = T;
type N = T::Num;
fn distance_to_aaline<A: Axis>(
&mut self,
point: Vec2<Self::N>,
axis: A,
val: Self::N,
) -> Self::N {
if axis.is_xaxis() {
(self.xline)(&mut self.acc, point, val)
} else {
(self.yline)(&mut self.acc, point, val)
}
}
fn distance_to_broad(
&mut self,
point: Vec2<Self::N>,
rect: PMut<Self::T>,
) -> Option<Self::N> {
(self.broad)(&mut self.acc, point, rect)
}
fn distance_to_fine(&mut self, point: Vec2<Self::N>, bot: PMut<Self::T>) -> Self::N {
(self.fine)(&mut self.acc, point, bot)
}
}
KnearestClosure {
_p: PhantomData,
acc,
broad,
fine,
xline,
yline,
}
}
#[derive(Debug)]
pub struct KnearestResult<'a, T: Aabb> {
pub bot: PMut<'a, T>,
pub mag: T::Num,
}
struct ClosestCand<'a, T: Aabb> {
bots: Vec<KnearestResult<'a, T>>,
curr_num: usize,
num: usize,
}
impl<'a, T: Aabb> ClosestCand<'a, T> {
fn into_sorted(self) -> Vec<KnearestResult<'a, T>> {
self.bots
}
fn new(num: usize) -> ClosestCand<'a, T> {
let bots = Vec::with_capacity(num);
ClosestCand {
bots,
num,
curr_num: 0,
}
}
fn consider<K: Knearest<T = T, N = T::Num>>(
&mut self,
point: &Vec2<K::N>,
knear: &mut K,
mut curr_bot: PMut<'a, T>,
) -> bool {
if let Some(long_dis) = knear.distance_to_broad(*point, curr_bot.borrow_mut()) {
if self.curr_num == self.num {
if let Some(l) = self.bots.last() {
if long_dis > l.mag {
return false;
}
}
}
}
let curr_dis = knear.distance_to_fine(*point, curr_bot.borrow_mut());
if self.curr_num < self.num {
let arr = &mut self.bots;
for i in 0..arr.len() {
if curr_dis < arr[i].mag {
let unit = KnearestResult {
bot: curr_bot,
mag: curr_dis,
};
arr.insert(i, unit);
self.curr_num += 1;
return true;
}
}
let unit = KnearestResult {
bot: curr_bot,
mag: curr_dis,
};
self.curr_num += 1;
arr.push(unit);
} else {
let arr = &mut self.bots;
for i in 0..arr.len() {
if curr_dis < arr[i].mag {
let v = arr.pop().unwrap();
loop {
if arr[arr.len() - 1].mag == v.mag {
arr.pop().unwrap();
} else {
break;
}
}
let unit = KnearestResult {
bot: curr_bot,
mag: curr_dis,
};
arr.insert(i, unit);
let max = arr
.iter()
.map(|a| a.mag)
.max_by(|a, b| {
if a > b {
Ordering::Greater
} else {
Ordering::Less
}
})
.unwrap();
assert!(max < v.mag);
return true;
} else if curr_dis == arr[i].mag {
let unit = KnearestResult {
bot: curr_bot,
mag: curr_dis,
};
arr.insert(i, unit);
return true;
}
}
}
false
}
fn full_and_max_distance(&self) -> Option<T::Num> {
use is_sorted::IsSorted;
assert!(IsSorted::is_sorted(&mut self.bots.iter().map(|a| a.mag)));
if self.curr_num == self.num {
self.bots.last().map(|a| a.mag)
} else {
None
}
}
}
struct Blap<'a, K: Knearest> {
knear: K,
point: Vec2<K::N>,
closest: ClosestCand<'a, K::T>,
}
impl<'a, K: Knearest> Blap<'a, K> {
fn should_recurse<A: Axis>(&mut self, line: (A, K::N)) -> bool {
if let Some(m) = self.closest.full_and_max_distance() {
let dis = self.knear.distance_to_aaline(self.point, line.0, line.1);
dis < m
} else {
true
}
}
}
fn recc<'a, 'b: 'a, T: Aabb, A: Axis, K: Knearest<N = T::Num, T = T>>(
axis: A,
stuff: LevelIter<VistrMut<'a, Node<'b, T>>>,
blap: &mut Blap<'a, K>,
) {
let ((_depth, nn), rest) = stuff.next();
let handle_node = match rest {
Some([left, right]) => {
let div = match nn.div {
Some(b) => b,
None => return,
};
let line = (axis, div);
if *blap.point.get_axis(axis) < div {
recc(axis.next(), left, blap);
if blap.should_recurse(line) {
recc(axis.next(), right, blap);
}
} else {
recc(axis.next(), right, blap);
if blap.should_recurse(line) {
recc(axis.next(), left, blap);
}
}
if !nn.range.is_empty() {
match nn.cont.contains_ext(*blap.point.get_axis(axis)) {
core::cmp::Ordering::Less => blap.should_recurse((axis, nn.cont.start)),
core::cmp::Ordering::Greater => blap.should_recurse((axis, nn.cont.end)),
core::cmp::Ordering::Equal => true,
}
} else {
false
}
}
None => true,
};
if handle_node {
for bot in nn.into_range().iter_mut() {
blap.closest.consider(&blap.point, &mut blap.knear, bot);
}
}
}
pub struct KResult<'a, T: Aabb> {
num_entires: usize,
inner: Vec<KnearestResult<'a, T>>,
}
impl<'a, T: Aabb> KResult<'a, T> {
#[inline(always)]
pub fn iter(
&mut self,
) -> impl Iterator<Item = &mut [KnearestResult<'a, T>]>
+ core::iter::FusedIterator
+ DoubleEndedIterator {
crate::util::SliceSplitMut::new(&mut self.inner, |a, b| a.mag == b.mag).fuse()
}
#[inline(always)]
pub fn into_vec(self) -> Vec<KnearestResult<'a, T>> {
self.inner
}
#[inline(always)]
pub fn total_len(&self) -> usize {
self.inner.len()
}
#[inline(always)]
pub fn len(&self) -> usize {
self.num_entires
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
}
pub fn assert_k_nearest_mut<T: Aabb>(
tree: &mut Tree<T>,
point: Vec2<T::Num>,
num: usize,
knear: &mut impl Knearest<T = T, N = T::Num>,
) {
let bots = tree.get_elements_mut();
use core::ops::Deref;
fn into_ptr_usize<T>(a: &T) -> usize {
a as *const T as usize
}
let mut res_naive = naive_k_nearest_mut(bots, point, num, knear)
.into_vec()
.drain(..)
.map(|a| (into_ptr_usize(a.bot.deref()), a.mag))
.collect::<Vec<_>>();
let r = tree.k_nearest_mut(point, num, knear);
let mut res_dino: Vec<_> = r
.into_vec()
.drain(..)
.map(|a| (into_ptr_usize(a.bot.deref()), a.mag))
.collect();
res_naive.sort_by(|a, b| a.partial_cmp(b).unwrap());
res_dino.sort_by(|a, b| a.partial_cmp(b).unwrap());
assert_eq!(res_naive.len(), res_dino.len());
assert!(res_naive.iter().eq(res_dino.iter()));
}
pub fn naive_k_nearest_mut<'a, T: Aabb>(
elems: PMut<'a, [T]>,
point: Vec2<T::Num>,
num: usize,
k: &mut impl Knearest<T = T, N = T::Num>,
) -> KResult<'a, T> {
let mut closest = ClosestCand::new(num);
for b in elems.iter_mut() {
closest.consider(&point, k, b);
}
let num_entires = closest.curr_num;
KResult {
num_entires,
inner: closest.into_sorted(),
}
}
use super::Queries;
pub trait KnearestQuery<'a>: Queries<'a> {
#[must_use]
fn k_nearest_mut<'b, K: Knearest<T = Self::T, N = Self::Num>>(
&'b mut self,
point: Vec2<Self::Num>,
num: usize,
ktrait: &mut K,
) -> KResult<Self::T>
where
'a: 'b,
{
let dt = self.vistr_mut().with_depth(Depth(0));
let knear = KnearestBorrow(ktrait);
let closest = ClosestCand::new(num);
let mut blap = Blap {
knear,
point,
closest,
};
recc(default_axis(), dt, &mut blap);
let num_entires = blap.closest.curr_num;
KResult {
num_entires,
inner: blap.closest.into_sorted(),
}
}
}