11use std:: { collections:: VecDeque , ops:: Range } ;
22
3+ pub struct Bitset < T : Copy > {
4+ curr : usize ,
5+ array : Vec < T > ,
6+ len : usize ,
7+ }
8+
9+ impl < T : Copy > Iterator for Bitset < T > {
10+ type Item = Vec < T > ;
11+
12+ fn next ( & mut self ) -> Option < Vec < T > > {
13+ if self . curr == ( 1 << self . len ) {
14+ return None ;
15+ }
16+
17+ let mut ret = Vec :: < T > :: new ( ) ;
18+ for ( i, & ai) in self . array . iter ( ) . enumerate ( ) {
19+ if ( self . curr >> i & 1 ) == 1 {
20+ ret. push ( ai) ;
21+ }
22+ }
23+
24+ self . curr += 1 ;
25+ Some ( ret)
26+ }
27+ }
28+
29+ pub fn bitset < T : Copy > ( a : Vec < T > ) -> Bitset < T > {
30+ let len = a. len ( ) ;
31+ Bitset {
32+ curr : 0 ,
33+ array : a,
34+ len,
35+ }
36+ }
37+
38+ #[ cfg( test) ]
39+ mod test_bitset {
40+ use crate :: collection:: bitset;
41+
42+ #[ test]
43+ fn it_works ( ) {
44+ let mut bitset = bitset ( vec ! [ 1 , 2 , 3 ] ) ;
45+ assert_eq ! ( bitset. next( ) , Some ( vec![ ] ) ) ;
46+ assert_eq ! ( bitset. next( ) , Some ( vec![ 1 ] ) ) ;
47+ assert_eq ! ( bitset. next( ) , Some ( vec![ 2 ] ) ) ;
48+ assert_eq ! ( bitset. next( ) , Some ( vec![ 1 , 2 ] ) ) ;
49+ assert_eq ! ( bitset. next( ) , Some ( vec![ 3 ] ) ) ;
50+ assert_eq ! ( bitset. next( ) , Some ( vec![ 1 , 3 ] ) ) ;
51+ assert_eq ! ( bitset. next( ) , Some ( vec![ 2 , 3 ] ) ) ;
52+ assert_eq ! ( bitset. next( ) , Some ( vec![ 1 , 2 , 3 ] ) ) ;
53+ assert ! ( bitset. next( ) . is_none( ) ) ;
54+ }
55+ }
56+
357#[ derive( Debug ) ]
458pub enum Item {
559 Pre ( usize ) ,
@@ -125,7 +179,7 @@ impl Iterator for CollectionIter<'_> {
125179
126180#[ cfg( test) ]
127181mod test_iterator {
128- use crate :: collection:: iterator :: CollectionIter ;
182+ use crate :: collection:: CollectionIter ;
129183 fn check ( iterator : CollectionIter , num_expected : usize , expected : Vec < Vec < usize > > ) {
130184 let mut num_count = 0 ;
131185 for perm in iterator {
@@ -136,7 +190,7 @@ mod test_iterator {
136190 }
137191
138192 mod with_duplication {
139- use crate :: collection:: iterator :: { test_iterator:: check, CollectionIter } ;
193+ use crate :: collection:: { test_iterator:: check, CollectionIter } ;
140194
141195 #[ test]
142196 fn it_works_permutation ( ) {
@@ -199,7 +253,7 @@ mod test_iterator {
199253 }
200254
201255 mod without_duplication {
202- use crate :: collection:: iterator :: { test_iterator:: check, CollectionIter } ;
256+ use crate :: collection:: { test_iterator:: check, CollectionIter } ;
203257
204258 #[ test]
205259 fn it_works_permutation ( ) {
@@ -246,3 +300,127 @@ mod test_iterator {
246300 }
247301 }
248302}
303+
304+ #[ macro_export]
305+ macro_rules! ndarray {
306+ // ndarray!(val; *shape)
307+ ( $x: expr; ) => { $x } ;
308+ ( $x: expr; $size: expr $( , $rest: expr ) * ) => {
309+ vec![ ndarray!( $x; $( $rest) ,* ) ; $size]
310+ } ;
311+ }
312+
313+ #[ cfg( test) ]
314+ mod test_ndarray {
315+
316+ #[ test]
317+ fn it_works ( ) {
318+ // ndarray!(val; 1) => [val]
319+ assert_eq ! ( ndarray!( 5 ; 1 ) , vec![ 5 ] ) ;
320+ // ndarray!(val; 1, 2) => [[val, val]]
321+ assert_eq ! ( ndarray!( 5 ; 1 , 2 ) , vec![ vec![ 5 , 5 ] ] ) ;
322+ // ndarray!(val; 2, 1) => [[val], [val]]
323+ assert_eq ! ( ndarray!( 5 ; 2 , 1 ) , vec![ vec![ 5 ] , vec![ 5 ] ] ) ;
324+ }
325+ }
326+
327+ #[ derive( Debug , Clone ) ]
328+ pub struct UnionFind {
329+ parents : Vec < usize > ,
330+ sizes : Vec < usize > ,
331+ }
332+
333+ #[ allow( clippy:: needless_range_loop) ]
334+ impl UnionFind {
335+ pub fn new ( n : usize ) -> Self {
336+ Self {
337+ parents : ( 0 ..n) . collect ( ) ,
338+ sizes : vec ! [ 1usize ; n] ,
339+ }
340+ }
341+
342+ pub fn parent ( & mut self , x : usize ) -> usize {
343+ if self . parents [ x] == x {
344+ x
345+ } else {
346+ self . parents [ x] = self . parent ( self . parents [ x] ) ;
347+ self . parents [ x]
348+ }
349+ }
350+
351+ pub fn unite ( & mut self , x : usize , y : usize ) {
352+ let mut px = self . parent ( x) ;
353+ let mut py = self . parent ( y) ;
354+
355+ if px == py {
356+ return ;
357+ }
358+
359+ if self . sizes [ px] < self . sizes [ py] {
360+ std:: mem:: swap ( & mut px, & mut py) ;
361+ }
362+
363+ self . sizes [ px] += self . sizes [ py] ;
364+ self . parents [ py] = px;
365+ }
366+
367+ pub fn size ( & mut self , x : usize ) -> usize {
368+ let x = self . parent ( x) ;
369+ self . sizes [ x]
370+ }
371+
372+ pub fn same ( & mut self , x : usize , y : usize ) -> bool {
373+ let px = self . parent ( x) ;
374+ let py = self . parent ( y) ;
375+ px == py
376+ }
377+ }
378+
379+ #[ cfg( test) ]
380+ mod test_union_find {
381+ use crate :: collection:: UnionFind ;
382+
383+ // helper function
384+ fn sizes ( uf : & mut UnionFind , n : usize ) -> Vec < usize > {
385+ ( 0 ..n) . map ( |i| uf. size ( i) ) . collect ( )
386+ }
387+
388+ #[ test]
389+ fn it_works ( ) {
390+ let n: usize = 5 ;
391+ let mut uf = UnionFind :: new ( n) ;
392+ assert_eq ! ( sizes( & mut uf, n) , [ 1 , 1 , 1 , 1 , 1 ] ) ;
393+
394+ uf. unite ( 0 , 1 ) ;
395+ assert_eq ! ( uf. parent( 0 ) , uf. parent( 1 ) ) ;
396+ assert ! ( uf. same( 0 , 1 ) ) ;
397+ assert_ne ! ( uf. parent( 0 ) , uf. parent( 2 ) ) ;
398+ assert ! ( !uf. same( 0 , 2 ) ) ;
399+ assert_eq ! ( sizes( & mut uf, n) , [ 2 , 2 , 1 , 1 , 1 ] ) ;
400+
401+ // check noop
402+ uf. unite ( 0 , 1 ) ;
403+ assert_eq ! ( uf. parent( 0 ) , uf. parent( 1 ) ) ;
404+ assert ! ( uf. same( 0 , 1 ) ) ;
405+ assert_ne ! ( uf. parent( 0 ) , uf. parent( 2 ) ) ;
406+ assert ! ( !uf. same( 0 , 2 ) ) ;
407+ assert_eq ! ( sizes( & mut uf, n) , [ 2 , 2 , 1 , 1 , 1 ] ) ;
408+
409+ uf. unite ( 0 , 2 ) ;
410+ assert_eq ! ( uf. parent( 0 ) , uf. parent( 2 ) ) ;
411+ assert ! ( uf. same( 0 , 2 ) ) ;
412+ assert_eq ! ( sizes( & mut uf, n) , [ 3 , 3 , 3 , 1 , 1 ] ) ;
413+
414+ uf. unite ( 3 , 4 ) ;
415+ assert_ne ! ( uf. parent( 0 ) , uf. parent( 3 ) ) ;
416+ assert ! ( !uf. same( 0 , 3 ) ) ;
417+ assert_eq ! ( sizes( & mut uf, n) , [ 3 , 3 , 3 , 2 , 2 ] ) ;
418+
419+ uf. unite ( 0 , 3 ) ;
420+ assert_eq ! ( uf. parent( 0 ) , uf. parent( 3 ) ) ;
421+ assert ! ( uf. same( 0 , 3 ) ) ;
422+ assert_eq ! ( uf. parent( 0 ) , uf. parent( 4 ) ) ;
423+ assert ! ( uf. same( 0 , 4 ) ) ;
424+ assert_eq ! ( sizes( & mut uf, n) , [ 5 , 5 , 5 , 5 , 5 ] ) ;
425+ }
426+ }
0 commit comments