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