1 ! Copyright (C) 2008 Eric Mertens.
2 ! See https://factorcode.org/license.txt for BSD license.
3 USING: accessors assocs hashtables kernel math sequences ;
7 { parents hashtable read-only }
8 { ranks hashtable read-only }
9 { counts hashtable read-only } ;
13 : add-count ( p a disjoint-set -- )
14 counts>> [ at '[ _ + ] ] [ swap change-at ] bi ; inline
16 : set-parent ( p a disjoint-set -- )
17 parents>> set-at ; inline
19 : link-sets ( p a disjoint-set -- )
20 [ set-parent ] [ add-count ] 3bi ; inline
22 : inc-rank ( a disjoint-set -- )
23 ranks>> [ 1 + ] change-at ; inline
27 GENERIC: representative ( a disjoint-set -- p )
29 M:: disjoint-set representative ( a disjoint-set -- p )
30 a disjoint-set parents>> at :> p
32 p disjoint-set representative [
33 a disjoint-set set-parent
39 : representatives ( a b disjoint-set -- r r )
40 '[ _ representative ] bi@ ; inline
42 : ranks ( a b disjoint-set -- r r )
43 '[ _ ranks>> at ] bi@ ; inline
45 :: branch ( a b neg zero pos -- )
46 a b = zero [ a b < neg pos if ] if ; inline
50 : <disjoint-set> ( -- disjoint-set )
51 H{ } clone H{ } clone H{ } clone disjoint-set boa ;
53 GENERIC: add-atom ( a disjoint-set -- )
55 M: disjoint-set add-atom
56 [ dupd parents>> set-at ]
57 [ [ 0 ] 2dip ranks>> set-at ]
58 [ [ 1 ] 2dip counts>> set-at ]
61 : add-atoms ( seq disjoint-set -- ) '[ _ add-atom ] each ;
63 GENERIC: disjoint-set-member? ( a disjoint-set -- ? )
65 M: disjoint-set disjoint-set-member? parents>> key? ;
67 GENERIC: disjoint-set-members ( disjoint-set -- seq )
69 M: disjoint-set disjoint-set-members parents>> keys ;
71 GENERIC: equiv-set-size ( a disjoint-set -- n )
73 M: disjoint-set equiv-set-size
74 [ representative ] keep counts>> at ;
76 GENERIC: equiv? ( a b disjoint-set -- ? )
78 M: disjoint-set equiv? representatives = ;
80 GENERIC: equate ( a b disjoint-set -- )
82 M:: disjoint-set equate ( a b disjoint-set -- )
83 a b disjoint-set representatives
85 2dup disjoint-set ranks
86 [ swap ] [ over disjoint-set inc-rank ] [ ] branch
87 disjoint-set link-sets
90 : equate-all-with ( seq a disjoint-set -- )
91 '[ _ _ equate ] each ;
93 : equate-all ( seq disjoint-set -- )
94 over empty? [ 2drop ] [
95 [ unclip-slice ] dip equate-all-with
99 [ parents>> ] [ ranks>> ] [ counts>> ] tri
100 [ clone ] tri@ disjoint-set boa ;
102 : assoc>disjoint-set ( assoc -- disjoint-set )
104 [ '[ drop _ add-atom ] assoc-each ]
105 [ '[ _ equate ] assoc-each ] 2bi