basis/disjoint-sets/disjoint-sets.factor

   1 ! Copyright (C) 2008 Eric Mertens.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors assocs hashtables kernel math sequences ;
   4 IN: disjoint-sets
   5
   6 TUPLE: disjoint-set
   7 { parents hashtable read-only }
   8 { ranks hashtable read-only }
   9 { counts hashtable read-only } ;
  10
  11 <PRIVATE
  12
  13 : add-count ( p a disjoint-set -- )
  14     counts>> [ at '[ _ + ] ] [ swap change-at ] bi ; inline
  15
  16 : set-parent ( p a disjoint-set -- )
  17     parents>> set-at ; inline
  18
  19 : link-sets ( p a disjoint-set -- )
  20     [ set-parent ] [ add-count ] 3bi ; inline
  21
  22 : inc-rank ( a disjoint-set -- )
  23     ranks>> [ 1 + ] change-at ; inline
  24
  25 PRIVATE>
  26
  27 GENERIC: representative ( a disjoint-set -- p )
  28
  29 M:: disjoint-set representative ( a disjoint-set -- p )
  30     a disjoint-set parents>> at :> p
  31     a p = [ a ] [
  32         p disjoint-set representative [
  33             a disjoint-set set-parent
  34         ] keep
  35     ] if ;
  36
  37 <PRIVATE
  38
  39 : representatives ( a b disjoint-set -- r r )
  40     '[ _ representative ] bi@ ; inline
  41
  42 : ranks ( a b disjoint-set -- r r )
  43     '[ _ ranks>> at ] bi@ ; inline
  44
  45 :: branch ( a b neg zero pos -- )
  46     a b = zero [ a b < neg pos if ] if ; inline
  47
  48 PRIVATE>
  49
  50 : <disjoint-set> ( -- disjoint-set )
  51     H{ } clone H{ } clone H{ } clone disjoint-set boa ;
  52
  53 GENERIC: add-atom ( a disjoint-set -- )
  54
  55 M: disjoint-set add-atom
  56     [ dupd parents>> set-at ]
  57     [ [ 0 ] 2dip ranks>> set-at ]
  58     [ [ 1 ] 2dip counts>> set-at ]
  59     2tri ;
  60
  61 : add-atoms ( seq disjoint-set -- ) '[ _ add-atom ] each ;
  62
  63 GENERIC: disjoint-set-member? ( a disjoint-set -- ? )
  64
  65 M: disjoint-set disjoint-set-member? parents>> key? ;
  66
  67 GENERIC: disjoint-set-members ( disjoint-set -- seq )
  68
  69 M: disjoint-set disjoint-set-members parents>> keys ;
  70
  71 GENERIC: equiv-set-size ( a disjoint-set -- n )
  72
  73 M: disjoint-set equiv-set-size
  74     [ representative ] keep counts>> at ;
  75
  76 GENERIC: equiv? ( a b disjoint-set -- ? )
  77
  78 M: disjoint-set equiv? representatives = ;
  79
  80 GENERIC: equate ( a b disjoint-set -- )
  81
  82 M:: disjoint-set equate ( a b disjoint-set -- )
  83     a b disjoint-set representatives
  84     2dup = [ 2drop ] [
  85         2dup disjoint-set ranks
  86         [ swap ] [ over disjoint-set inc-rank ] [ ] branch
  87         disjoint-set link-sets
  88     ] if ;
  89
  90 : equate-all-with ( seq a disjoint-set -- )
  91     '[ _ _ equate ] each ;
  92
  93 : equate-all ( seq disjoint-set -- )
  94     over empty? [ 2drop ] [
  95         [ unclip-slice ] dip equate-all-with
  96     ] if ;
  97
  98 M: disjoint-set clone
  99     [ parents>> ] [ ranks>> ] [ counts>> ] tri
 100     [ clone ] tri@ disjoint-set boa ;
 101
 102 : assoc>disjoint-set ( assoc -- disjoint-set )
 103     <disjoint-set> [
 104         [ '[ drop _ add-atom ] assoc-each ]
 105         [ '[ _ equate ] assoc-each ] 2bi
 106     ] keep ;