1use std::borrow::Borrow;
2use std::cell::RefCell;
3use std::fmt::Debug;
4use std::marker::PhantomData;
5use std::ops::{Add, Deref, Mul, Neg, Sub};
6use std::rc::Rc;
7
8#[derive(Debug,Clone,Copy,PartialEq,Eq)]
9pub enum Color {
10 Red,
11 Black
12}
13#[derive(Debug,Clone,Copy,PartialEq,Eq)]
14pub enum LR {
15 L,
16 R
17}
18#[derive(Debug,Clone,Copy,PartialEq,Eq)]
19pub enum Balance {
20 Pre,
21 Fix,
22 None
23}
24impl Neg for Color {
25 type Output = Color;
26 fn neg(self) -> Self::Output {
27 match self {
28 Color::Red => Color::Black,
29 Color::Black => Color::Red
30 }
31 }
32}
33pub trait Square: Mul + Sized + Clone + Copy {
34 fn square(self) -> <Self as Mul>::Output {
35 self * self
36 }
37}
38impl<T> Square for T where T: Mul + Sized + Clone + Copy {}
39pub trait EuclideanDistance<T> {
40 type Output;
41 fn euclidean_distance(self, rhs:T) -> Self::Output;
42}
43impl<'a,const K:usize,P> EuclideanDistance<&'a [P; K]> for &'a [P; K]
44 where P: PartialOrd + Mul<P, Output = P> + Add<P, Output = P> + Sub<P, Output = P> +
45 Clone + Copy + Default + Distance<P, Output = P> + Square + Sized {
46 type Output = P;
47
48 fn euclidean_distance(self, rhs: &'a [P; K]) -> P {
49 self.iter().zip(rhs.iter()).fold(P::default(),|acc,(p1,p2)| {
50 acc + p1.distance(p2).square()
51 })
52 }
53}
54pub trait Distance<T> {
55 type Output;
56
57 fn distance(&self, rhs:&T) -> Self::Output;
58}
59impl<P> Distance<P> for P
60 where P: PartialOrd + Mul<P, Output = P> + Add<P, Output = P> + Sub<P, Output = P> +
61 Clone + Copy + Default {
62 type Output = P;
63
64 fn distance(&self, rhs: &P) -> P {
65 if self.partial_cmp(rhs).unwrap().is_le() {
66 *rhs - *self
67 } else {
68 *self - *rhs
69 }
70 }
71}
72#[derive(Debug)]
73pub struct KDNode<'a,const K:usize,P,T>
74 where P: Debug + PartialOrd + Mul<Output = P> + Add + Sub +
75 Clone + Copy + Default + Distance<P, Output = P> + Square + Sized + 'a,
76 &'a [P; K]: EuclideanDistance<&'a [P; K], Output = P> + 'a {
77 positions:Rc<[P; K]>,
78 value:Rc<RefCell<T>>,
79 color: Rc<RefCell<Color>>,
80 left:Option<Box<KDNode<'a, K,P,T>>>,
81 right:Option<Box<KDNode<'a, K,P,T>>>,
82 demension:usize,
83 l:PhantomData<&'a ()>
84}
85impl<'a,const K:usize,P,T> KDNode<'a, K,P,T>
86 where P: Debug + PartialOrd + Mul<Output = P> + Add + Sub +
87 Clone + Copy + Default + Distance<P, Output = P> + 'a,
88 &'a [P; K]: EuclideanDistance<&'a [P; K], Output = P> + 'a {
89 pub fn new(positions:Rc<[P; K]>, value:Rc<RefCell<T>>, demension:usize) -> KDNode<'a, K,P,T> {
90 KDNode {
91 positions: positions,
92 value: value,
93 color: Rc::new(RefCell::new(Color::Red)),
94 left: None,
95 right: None,
96 demension: demension,
97 l:PhantomData::<&'a ()>
98 }
99 }
100
101 fn with_color(positions:Rc<[P; K]>, value:Rc<RefCell<T>>, color:Rc<RefCell<Color>>, demension:usize) -> KDNode<'a, K,P,T> {
102 KDNode {
103 positions: positions,
104 value: value,
105 color: color,
106 left: None,
107 right: None,
108 demension: demension,
109 l:PhantomData::<&'a ()>
110 }
111 }
112
113 pub fn right_rotate(t: KDNode<'a, K,P,T>) -> KDNode<'a, K,P,T> {
114 match t.left {
115 Some(left) => {
116 KDNode {
117 positions: left.positions,
118 value: left.value,
119 color: left.color,
120 left: left.left,
121 right: Some(Box::new(KDNode {
122 positions: t.positions,
123 value: t.value,
124 color: t.color,
125 left: left.right,
126 right: t.right,
127 demension: t.demension,
128 l:PhantomData::<&'a ()>
129 },)),
130 demension: left.demension,
131 l:PhantomData::<&'a ()>
132 }
133 },
134 None => t
135 }
136 }
137
138 pub fn left_rotate(t: KDNode<'a, K,P,T>) -> KDNode<'a, K,P,T> {
139 match t.right {
140 Some(right) => {
141 KDNode {
142 positions: right.positions,
143 value: right.value,
144 color: right.color,
145 right: right.right,
146 left: Some(Box::new(KDNode {
147 positions: t.positions,
148 value: t.value,
149 color: t.color,
150 right: right.left,
151 left: t.left,
152 demension: t.demension,
153 l:PhantomData::<&'a ()>
154 })),
155 demension: right.demension,
156 l:PhantomData::<&'a ()>
157 }
158 },
159 None => t
160 }
161 }
162
163 #[allow(dead_code)]
164 fn left_and_right_rotate(mut t: KDNode<'a, K,P,T>) -> KDNode<'a, K,P,T> {
165 match t.left.take() {
166 None => {
167 t
168 },
169 Some(left) => {
170 t.left = Some(Box::new(Self::left_rotate(*left)));
171 Self::right_rotate(t)
172 }
173 }
174 }
175
176 #[allow(dead_code)]
177 fn right_and_left_rotate(mut t: KDNode<'a, K,P,T>) -> KDNode<'a, K,P,T> {
178 match t.right.take() {
179 None => {
180 t
181 },
182 Some(right) => {
183 t.right = Some(Box::new(Self::right_rotate(*right)));
184 Self::left_rotate(t)
185 }
186 }
187 }
188
189 fn nearest(t: Option<&'a Box<KDNode<'a, K,P,T>>>,
190 positions:&'a [P; K],
191 distance:P,
192 nearest_positions:&'a [P; K],
193 current_value:&Rc<RefCell<T>>,
194 demension:usize) -> Option<(P, &'a [P; K], Rc<RefCell<T>>)> {
195 t.and_then(|t| {
196 if positions[demension].partial_cmp(&t.positions[demension]).unwrap().is_lt() {
197 if let Some(c) = t.left.as_ref() {
198 if let Some((distance,current_positions,current_value)) = Self::nearest(
199 Some(&c), positions, distance, nearest_positions, current_value, (demension + 1) % K) {
200 let (distance,current_positions,current_value) = Self::nearest_center(Some(t),
201 positions,
202 distance,
203 current_positions,
204 ¤t_value,
205 demension).unwrap();
206
207 if let Some(c) = t.right.as_ref() {
208 if distance.partial_cmp(&positions[demension].distance(&c.positions[demension]).square()).unwrap().is_lt() {
209 Some((distance,current_positions,current_value))
210 } else {
211 Self::nearest(Some(&c),positions,distance,current_positions,¤t_value,(demension + 1) % K)
212 }
213 } else {
214 Some((distance,current_positions,current_value))
215 }
216 } else {
217 unreachable!()
218 }
219 } else {
220 let (distance,current_positions,current_value) = Self::nearest_center(Some(t),
221 positions,
222 distance,
223 nearest_positions,
224 ¤t_value,
225 demension).unwrap();
226
227 if let Some(c) = t.right.as_ref() {
228 if distance.partial_cmp(&positions[demension].distance(&c.positions[demension]).square()).unwrap().is_lt() {
229 Some((distance,current_positions,current_value))
230 } else {
231 Self::nearest(Some(&c),positions,distance,current_positions,¤t_value,(demension + 1) % K)
232 }
233 } else {
234 Some((distance,current_positions,current_value))
235 }
236 }
237 } else {
238 if let Some(c) = t.right.as_ref() {
239 if let Some((distance,current_positions,current_value)) = Self::nearest(
240 Some(&c),positions,distance,nearest_positions,current_value,(demension + 1) % K) {
241 let (distance,current_positions,current_value) = Self::nearest_center(Some(t),
242 positions,
243 distance,
244 current_positions,
245 ¤t_value,
246 demension).unwrap();
247
248 if let Some(c) = t.left.as_ref() {
249 if distance.partial_cmp(&positions[demension].distance(&c.positions[demension]).square()).unwrap().is_lt() {
250 Some((distance,current_positions,current_value))
251 } else {
252 Self::nearest(Some(&c),positions,distance,current_positions,¤t_value,(demension + 1) % K)
253 }
254 } else {
255 Some((distance,current_positions,current_value))
256 }
257 } else {
258 unreachable!()
259 }
260 } else {
261 let (distance,current_positions,current_value) = Self::nearest_center(Some(t),
262 positions,
263 distance,
264 nearest_positions,
265 ¤t_value,
266 demension).unwrap();
267
268
269 if let Some(c) = t.left.as_ref() {
270 if distance.partial_cmp(&positions[demension].distance(&c.positions[demension]).square()).unwrap().is_lt() {
271 Some((distance,current_positions,current_value))
272 } else {
273 Self::nearest(Some(&c),positions,distance,current_positions,¤t_value,(demension + 1) % K)
274 }
275 } else {
276 Some((distance,current_positions,current_value))
277 }
278 }
279 }
280 })
281 }
282
283 fn nearest_center(t: Option<&'a Box<KDNode<'a, K,P,T>>>,
284 positions:&'a [P; K],
285 mut distance:P,
286 nearest_positions:&'a [P; K],
287 current_value:&Rc<RefCell<T>>,
288 _:usize) -> Option<(P, &'a [P; K], Rc<RefCell<T>>)> {
289 Some(t.and_then(|t| {
290 let d = positions.euclidean_distance(&t.positions);
291
292 let mut current_value = Rc::clone(¤t_value);
293
294 let current_positions = if d.partial_cmp(&distance).unwrap().is_le() {
295 distance = d;
296 current_value = Rc::clone(&t.value);
297 t.positions.borrow()
298 } else {
299 nearest_positions
300 };
301
302 Some((distance,current_positions, current_value))
303 }).expect("current node is none."))
304 }
305
306 fn insert(t: Option<Box<KDNode<'a, K,P,T>>>,
307 positions:&Rc<[P; K]>,
308 color:&Rc<RefCell<Color>>,
309 parent_color:Option<Color>,
310 lr:Option<LR>,
311 value:Rc<RefCell<T>>,
312 demension:usize) -> (KDNode<'a, K,P,T>, Balance) {
313 match t {
314 None if demension == K - 1 => {
315 let b = if parent_color.map(|c| c == Color::Red).unwrap_or(false) {
316 Balance::Pre
317 } else {
318 Balance::None
319 };
320
321 let color = if let Some(_) = parent_color {
322 Rc::new(RefCell::new(Color::Black))
323 } else {
324 Rc::clone(&color)
325 };
326
327 (KDNode::with_color(Rc::clone(positions), Rc::clone(&value),color,demension),b)
328 },
329 None if demension == 0 => {
330 let t = KDNode {
331 positions: Rc::clone(positions),
332 value: Rc::clone(&value),
333 color: Rc::clone(color),
334 left: None,
335 right: None,
336 demension: demension,
337 l:PhantomData::<&'a ()>
338 };
339
340 (t,Balance::None)
341 },
342 None => {
343 let t = KDNode {
344 positions: Rc::clone(positions),
345 value: Rc::clone(&value),
346 color: Rc::clone(&color),
347 left: None,
348 right: None,
349 demension: demension,
350 l:PhantomData::<&'a ()>
351 };
352
353 (t,Balance::None)
354 },
355 Some(mut t) if demension == K - 1 => {
356 let parent_color = Some(color.deref().borrow().clone());
357
358 if positions[demension].partial_cmp(&t.positions[demension]).unwrap().is_lt() {
359 let (n,b) = Self::insert(t.left,
360 positions,
361 &Rc::clone(&t.color),
362 parent_color,
363 Some(LR::L),
364 value, (demension+1) % K);
365
366 t.left = Some(Box::new(n));
367
368 (*t,b)
369 } else {
370 let (n,b) = Self::insert(t.right,
371 positions,
372 &Rc::clone(&t.color),
373 parent_color,
374 Some(LR::R),
375 value, (demension+1) % K);
376
377 t.right = Some(Box::new(n));
378
379 (*t,b)
380 }
381 },
382 Some(mut t) if demension == 0 => {
383 if positions[demension].partial_cmp(&t.positions[demension]).unwrap().is_lt() {
384 let (n,b) = Self::insert(t.left,
385 positions,
386 color,
387 parent_color,
388 lr,
389 value, (demension+1) % K);
390
391 t.left = Some(Box::new(n));
392
393 Self::balance(*t, demension, b, lr,Some(LR::L))
394 } else {
395 let (n,b) = Self::insert(t.right,
396 positions,
397 color,
398 parent_color,
399 lr,
400 value, (demension+1) % K);
401
402 t.right = Some(Box::new(n));
403
404 Self::balance(*t, demension, b, lr,Some(LR::R))
405 }
406 },
407 Some(mut t) => {
408 if positions[demension].partial_cmp(&t.positions[demension]).unwrap().is_lt() {
409 let (n,b) = Self::insert(t.left,
410 positions,
411 color,
412 parent_color,
413 lr,
414 value, (demension+1) % K);
415
416 t.left = Some(Box::new(n));
417
418 (*t,b)
419 } else {
420 let (n,b) = Self::insert(t.right,
421 positions,
422 color,
423 parent_color,
424 lr,
425 value, (demension+1) % K);
426
427 t.right = Some(Box::new(n));
428
429 (*t,b)
430 }
431 }
432 }
433 }
434
435 fn balance(mut t: KDNode<'a, K,P,T>, demension:usize, balance:Balance, parent_lr:Option<LR>, lr:Option<LR>) -> (KDNode<'a, K,P,T>, Balance) {
436 if demension > 0 {
437 (t,balance)
438 } else {
439 match balance {
440 Balance::None => (t, balance),
441 Balance::Pre => {
442 let lr = lr.unwrap();
443 let parent_lr = parent_lr.unwrap();
444
445 for _ in 0..K {
446 t = if parent_lr != lr && lr == LR::L {
447 Self::right_rotate(t)
448 } else if parent_lr != lr && lr == LR::R {
449 Self::left_rotate(t)
450 } else {
451 t
452 };
453 }
454
455 (t, Balance::Fix)
456 },
457 Balance::Fix => {
458 let lr = lr.unwrap();
459
460 for _ in 0..K {
461 t = match lr {
462 LR::L => Self::right_rotate(t),
463 LR::R => Self::left_rotate(t)
464 };
465 }
466
467 match lr {
468 LR::L => {
469 if let Some(c) = t.left.as_ref() {
470 *c.color.borrow_mut() = Color::Black;
471 }
472 },
473 LR::R => {
474 if let Some(c) = t.right.as_ref() {
475 *c.color.borrow_mut() = Color::Black;
476 }
477 }
478 }
479 (t, Balance::None)
480 }
481 }
482 }
483 }
484}
485#[derive(Debug)]
486pub struct KDTree<'a,const K:usize,P,T>
487 where P: Debug + PartialOrd + Mul<Output = P> + Add + Sub + Clone + Copy + Default + Distance<P, Output = P> + Square + Sized + 'a,
488 &'a [P; K]: EuclideanDistance<&'a [P; K], Output = P> + 'a {
489 root: Option<Box<KDNode<'a, K,P,T>>>,
490 l:PhantomData<&'a ()>
491}
492impl<'a,const K:usize,P,T> KDTree<'a, K,P,T>
493 where P: Debug + PartialOrd + Mul<Output = P> + Add + Sub + Clone + Copy + Default + Distance<P, Output = P> + 'a,
494 &'a [P; K]: EuclideanDistance<&'a [P; K], Output = P> + 'a {
495 pub fn new() -> KDTree<'a, K,P,T> {
496 KDTree {
497 root: None,
498 l:PhantomData::<&'a ()>
499 }
500 }
501
502 pub fn nearest(&'a self, positions:&'a [P; K]) -> Option<(&'a [P; K], Rc<RefCell<T>>)> {
503 self.root.as_ref().and_then(|root| {
504 let distance = positions.euclidean_distance(&root.positions);
505
506 KDNode::nearest(Some(root),positions,distance,&root.positions,&root.value,0).map(|(_,p,v)| {
507 (p,v)
508 })
509 })
510 }
511
512 pub fn nearest_position(&'a self, positions:&'a [P; K]) -> Option<&'a [P; K]> {
513 self.root.as_ref().and_then(|root| {
514 let distance = positions.euclidean_distance(&root.positions);
515
516 KDNode::nearest(Some(root),positions,distance,&root.positions,&root.value,0).map(|(_,p,_)| {
517 p
518 })
519 })
520 }
521
522 pub fn insert(&mut self, positions:[P; K], value:T) {
523 let (n,_) = KDNode::insert(self.root.take(),
524 &Rc::new(positions),
525 &Rc::new(RefCell::new(Color::Black)),
526 None,
527 None,
528 Rc::new(RefCell::new(value)),
529 0);
530 self.root = Some(Box::new(n));
531 self.root.as_ref().map(|root| {
532 *root.color.borrow_mut() = Color::Black;
533 });
534 }
535}
536
537