b_k_tree/
set.rs

1impl<'a,E,M:DiscreteMetric<E,E>+DiscreteMetric<E,Q>,Q> Iterator for DrainSetIter<'a,E,M,Q>{
2	fn next(&mut self)->Option<Self::Item>{self.inner.next().map(|(k,_v,d)|(k,d))}
3	fn size_hint(&self)->(usize,Option<usize>){self.inner.size_hint()}
4	type Item=(E,usize);
5}
6impl<'a,E,M:DiscreteMetric<E,Q>,Q> Iterator for CloseSetIter<'a,E,M,Q>{
7	fn next(&mut self)->Option<Self::Item>{self.inner.next()}
8	fn size_hint(&self)->(usize,Option<usize>){self.inner.size_hint()}
9	type Item=(&'a E,usize);
10}
11impl<'a,E,M,Q:Clone> Clone for CloseSetIter<'a,E,M,Q>{
12	fn clone(&self)->Self{
13		Self{inner:self.inner.clone()}
14	}
15	fn clone_from(&mut self,other:&Self){self.inner.clone_from(&other.inner)}
16}
17impl<'a,E,M> IntoIterator for &'a BKTreeSet<E,M>{
18	fn into_iter(self)->Self::IntoIter{
19		SetIter{inner:self.inner.keys()}
20	}
21	type IntoIter=SetIter<'a,E>;
22	type Item=&'a E;
23}
24impl<'a,E> Clone for SetIter<'a,E>{
25	fn clone(&self)->Self{
26		Self{inner:self.inner.clone()}
27	}
28	fn clone_from(&mut self,other:&Self){self.inner.clone_from(&other.inner)}
29}
30impl<'a,E> ExactSizeIterator for SetIter<'a,E>{
31	fn len(&self)->usize{self.inner.len()}
32}
33impl<'a,E> Iterator for SetIter<'a,E>{
34	fn next(&mut self)->Option<Self::Item>{self.inner.next()}
35	fn size_hint(&self)->(usize,Option<usize>){self.inner.size_hint()}
36	type Item=&'a E;
37}
38impl<E:Clone> Clone for SetIntoIter<E>{
39	fn clone(&self)->Self{
40		Self{inner:self.inner.clone()}
41	}
42	fn clone_from(&mut self,other:&Self){self.inner.clone_from(&other.inner)}
43}
44impl<E,M:Default> Default for BKTreeSet<E,M>{
45	fn default()->Self{Self::new(M::default())}
46}
47impl<E,M:Default+DiscreteMetric<E,E>> FromIterator<E> for BKTreeSet<E,M>{
48	fn from_iter<I:IntoIterator<Item=E>>(iter:I)->Self{
49		Self{inner:iter.into_iter().map(|e|(e,())).collect()}
50	}
51}
52impl<E,M:DiscreteMetric<E,E>> Extend<E> for BKTreeSet<E,M>{
53	fn extend<I:IntoIterator<Item=E>>(&mut self,iter:I){iter.into_iter().map(|e|self.insert(e)).for_each(|_|())}
54}
55impl<E,M> AsMut<Self> for BKTreeSet<E,M>{
56	fn as_mut(&mut self)->&mut Self{self}
57}
58impl<E,M> AsRef<Self> for BKTreeSet<E,M>{
59	fn as_ref(&self)->&Self{self}
60}
61impl<E,M> BKTreeSet<E,M>{//TODO other sterotypical set operations
62	/// moves all elements from other into self, leaving other empty
63	pub fn append<M2>(&mut self,other:&mut BKTreeSet<E,M2>) where M:DiscreteMetric<E,E>{self.inner.append(&mut other.inner)}
64	/// gets the items whose distance is at most max distance from key
65	pub fn close_iter<Q>(&self,key:Q,maxdistance:usize)->CloseSetIter<'_,E,M,Q> where M:DiscreteMetric<E,Q>{
66		CloseSetIter{inner:self.inner.close_keys(key,maxdistance)}
67	}
68	/// returns the elements at most maxdistance from the key, sorted by distance
69	pub fn close_sorted<'a,Q:?Sized>(&self,key:&Q,maxdistance:usize)->Vec<(&E,usize)> where M:DiscreteMetric<E,Q>{self.inner.close_sorted(key,maxdistance).into_iter().map(|(k,_v,d)|(k,d)).collect()}
70	/// tests if the set contains an element within max distance of the key
71	pub fn contains<Q:?Sized>(&self,key:&Q,maxdistance:usize)->bool where M:DiscreteMetric<E,Q>{self.inner.contains_key(key,maxdistance)}
72	/// drains the elements close to the key
73	pub fn drain<Q>(&mut self,key:Q,maxdistance:usize)->DrainSetIter<'_,E,M,Q> where M:DiscreteMetric<E,E>+DiscreteMetric<E,Q>{
74		DrainSetIter{inner:self.inner.drain(key,maxdistance)}
75	}
76	/// returns a reference to the element in the set that is closest to the key within max distance, or None if the set contains no element at most max distance from the given element. If there are multiple closest elements, exactly which is returned is unspecified
77	pub fn get<Q:?Sized>(&self,key:&Q,maxdistance:usize)->Option<(&E,usize)> where M:DiscreteMetric<E,Q>{self.inner.get_key_value(key,maxdistance).map(|(k,_v,d)|(k,d))}
78	/// inserts
79	pub fn insert(&mut self,value:E)->bool where M:DiscreteMetric<E,E>{self.inner.insert(value,()).is_none()}
80	/// returns true if the set contains no elements
81	pub fn is_empty(&self)->bool{self.inner.is_empty()}
82	/// makes an iterator over the items
83	pub fn iter(&self)->SetIter<'_,E>{
84		SetIter{inner:self.inner.keys()}
85	}
86	/// returns the number of elements in the set
87	pub fn len(&self)->usize{self.inner.len()}
88	/// references the metric. avoid modifying the metric in a way that changes the distances because that will most likely cause unspecified incorrect behavior
89	pub fn metric(&self)->&M{self.inner.metric()}
90	/// creates a new tree
91	pub fn new(metric:M)->Self{
92		Self{inner:BKTreeMap::new(metric)}
93	}
94	/// removes an item from the tree. This particular tree type doesn't allow super efficient removal, so try to avoid using too much.
95	pub fn remove<Q:?Sized>(&mut self,key:&Q,maxdistance:usize)->bool where M:DiscreteMetric<E,Q>+DiscreteMetric<E,E>{self.inner.remove_entry(key,maxdistance).is_some()}
96	/// removes all the elements for which f returns false
97	pub fn retain<F:FnMut(&E)->bool>(&mut self,mut f:F) where M:DiscreteMetric<E,E>{self.inner.retain(|k,_v|f(k))}
98	/// removes an item from the tree. This particular tree type doesn't allow super efficient removal, so try to avoid using too much.
99	pub fn take<Q:?Sized>(&mut self,key:&Q,maxdistance:usize)->Option<(E,usize)> where M:DiscreteMetric<E,Q>+DiscreteMetric<E,E>{self.inner.remove_entry(key,maxdistance).map(|(k,_v,d)|(k,d))}
100}
101impl<E,M> IntoIterator for BKTreeSet<E,M>{
102	fn into_iter(self)->Self::IntoIter{
103		SetIntoIter{inner:self.inner.into_keys()}
104	}
105	type IntoIter=SetIntoIter<E>;
106	type Item=E;
107}
108impl<E> ExactSizeIterator for SetIntoIter<E>{
109	fn len(&self)->usize{self.inner.len()}
110}
111impl<E> Iterator for SetIntoIter<E>{
112	fn next(&mut self)->Option<Self::Item>{self.inner.next()}
113	fn size_hint(&self)->(usize,Option<usize>){self.inner.size_hint()}
114	type Item=E;
115}
116#[cfg(test)]
117mod tests{
118	#[test]
119	fn contains_4(){
120		let mut tree=BKTreeSet::new(Cheb2D);
121
122		assert!(tree.insert((3,3)));
123		assert!(tree.contains(&(4,4),1));
124		assert!(tree.insert((3,4)));
125		assert!(tree.insert((3,5)));
126		assert!(tree.insert((4,3)));
127		assert!(tree.insert((4,5)));
128		assert!(tree.contains(&(4,4),1));
129		assert!(tree.insert((5,3)));
130		assert!(tree.insert((5,4)));
131		assert!(tree.insert((5,5)));
132
133		assert!(!tree.contains(&(4,4),0));
134		assert!(tree.contains(&(4,4),1));
135		assert!(tree.contains(&(4,4),2));
136		assert!(tree.contains(&(3,4),0));
137		assert!(tree.contains(&(4,3),0));
138		assert!(tree.contains(&(4,5),0));
139		assert!(tree.contains(&(5,4),0));
140		assert!(!tree.contains(&(5,2),0));
141		assert!(!tree.contains(&(5,6),0));
142		assert!(!tree.contains(&(6,5),0));
143		assert!(tree.contains(&(5,2),1));
144		assert!(tree.contains(&(5,6),1));
145		assert!(tree.contains(&(6,5),1));
146
147		assert_eq!(tree.take(&(5,5),0),Some(((5,5),0)));
148		assert_eq!(tree.take(&(5,3),2),Some(((5,3),0)));
149		assert!(tree.take(&(5,5),1).is_some());
150		assert!(tree.take(&(5,5),1).is_some());
151		assert!(tree.take(&(5,5),1).is_none());
152		assert!(!tree.contains(&(5,5),0));
153		assert!(tree.remove(&(4,4),1));
154		assert!(!tree.contains(&(4,4),0));
155		assert!(tree.contains(&(4,4),1));
156	}
157	#[test]
158	fn retain_cluster(){
159		let mut tree=BKTreeSet::new(Cheb2D);
160
161		tree.insert((0,0));
162		tree.insert((1,0));
163		tree.insert((0,1));
164		tree.insert((1,1));
165
166		tree.insert((10,10));
167		tree.insert((11,10));
168		tree.insert((10,11));
169		tree.insert((11,11));
170
171		tree.retain(|&(x,y)|x<10&&y<10);
172
173		let mut a:Vec<(isize,isize,usize)>=tree.close_iter((-1,-1),2).map(|((x,y),d)|(*x,*y,d)).collect();
174		let mut b:Vec<(isize,isize,usize)>=tree.close_iter((12,12),2).map(|((x,y),d)|(*x,*y,d)).collect();
175		a.sort_unstable();
176		b.sort_unstable();
177		assert_eq!(a,[(0,0,1),(0,1,2),(1,0,2),(1,1,2)]);
178		assert_eq!(b,[]);
179	}
180	#[test]
181	fn two_clusters(){
182		let mut tree=BKTreeSet::new(Cheb2D);
183
184		tree.insert((0,0));
185		tree.insert((1,0));
186		tree.insert((0,1));
187		tree.insert((1,1));
188
189		tree.insert((10,10));
190		tree.insert((11,10));
191		tree.insert((10,11));
192		tree.insert((11,11));
193
194		let mut a:Vec<(isize,isize,usize)>=tree.close_iter((-1,-1),2).map(|((x,y),d)|(*x,*y,d)).collect();
195		let mut b:Vec<(isize,isize,usize)>=tree.close_iter((12,12),2).map(|((x,y),d)|(*x,*y,d)).collect();
196		a.sort_unstable();
197		b.sort_unstable();
198		assert_eq!(a,[(0,0,1),(0,1,2),(1,0,2),(1,1,2)]);
199		assert_eq!(b,[(10,10,2),(10,11,2),(11,10,2),(11,11,1)]);
200	}
201	impl DiscreteMetric<(isize,isize),(isize,isize)> for Cheb2D{
202		fn distance(&self,x:&(isize,isize),y:&(isize,isize))->usize{
203			let ((xx,xy),(yx,yy))=(x,y);
204			((xx-yx).abs() as usize).max((xy-yy).abs() as usize)
205		}
206	}
207	#[derive(Clone,Copy,Debug,Default)]
208	/// 2d integer chebyshev distance
209	struct Cheb2D;
210	use super::*;
211}
212#[derive(Clone,Debug)]
213/// a set for quickly finding items separated by a small discrete distance, implemented as a thin wrapper over BKTreeMap
214pub struct BKTreeSet<E,M>{inner:BKTreeMap<E,M,()>}
215#[derive(Debug)]
216/// iterator over the items close to some key
217pub struct CloseSetIter<'a,E,M,Q>{inner:CloseKeyIter<'a,E,M,Q,()>}
218#[derive(Debug)]
219/// draining iterator over the set
220pub struct DrainSetIter<'a,E,M:DiscreteMetric<E,E>+DiscreteMetric<E,Q>,Q>{inner:DrainMapIter<'a,E,M,Q,()>}
221#[derive(Debug)]
222/// iterator over the items in the tree
223pub struct SetIntoIter<E>{inner:IntoKeysIter<E,()>}
224#[derive(Debug)]
225/// iterator over the items in the tree
226pub struct SetIter<'a,E>{inner:KeyIter<'a,E,()>}
227use {
228	crate::{
229		DiscreteMetric,map::{BKTreeMap,CloseKeyIter,DrainMapIter,IntoKeysIter,KeyIter}
230	},
231	std::iter::{Extend,FromIterator}
232};