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 GENERIC: representative ( a disjoint-set -- p )
  39
  40 M: disjoint-set representative
  41     2dup representative? [ drop ] [
  42         [ [ parent ] keep representative dup ] 2keep set-parent
  43     ] if ;
  44
  45 : representatives ( a b disjoint-set -- r r )
  46     [ representative ] curry bi@ ; inline
  47
  48 : ranks ( a b disjoint-set -- r r )
  49     [ rank ] curry bi@ ; inline
  50
  51 :: branch ( a b neg zero pos -- )
  52     a b = zero [ a b < neg pos if ] if ; inline
  53
  54 PRIVATE>
  55
  56 : <disjoint-set> ( -- disjoint-set )
  57     H{ } clone H{ } clone H{ } clone disjoint-set boa ;
  58
  59 GENERIC: add-atom ( a disjoint-set -- )
  60
  61 M: disjoint-set add-atom
  62     [ dupd parents>> set-at ]
  63     [ [ 0 ] 2dip ranks>> set-at ]
  64     [ [ 1 ] 2dip counts>> set-at ]
  65     2tri ;
  66
  67 : add-atoms ( seq disjoint-set -- ) '[ _ add-atom ] each ;
  68
  69 GENERIC: disjoint-set-member? ( a disjoint-set -- ? )
  70
  71 M: disjoint-set disjoint-set-member? parents>> key? ;
  72
  73 GENERIC: equiv-set-size ( a disjoint-set -- n )
  74
  75 M: disjoint-set equiv-set-size [ representative ] keep count ;
  76
  77 GENERIC: equiv? ( a b disjoint-set -- ? )
  78
  79 M: disjoint-set equiv? representatives = ;
  80
  81 GENERIC: equate ( a b disjoint-set -- )
  82
  83 M:: disjoint-set equate ( a b disjoint-set -- )
  84     a b disjoint-set representatives
  85     2dup = [ 2drop ] [
  86         2dup disjoint-set ranks
  87         [ swap ] [ over disjoint-set inc-rank ] [ ] branch
  88         disjoint-set link-sets
  89     ] if ;
  90
  91 : equate-all-with ( seq a disjoint-set -- )
  92     '[ _ _ equate ] each ;
  93
  94 : equate-all ( seq disjoint-set -- )
  95     over empty? [ 2drop ] [
  96         [ unclip-slice ] dip equate-all-with
  97     ] if ;
  98
  99 M: disjoint-set clone
 100     [ parents>> ] [ ranks>> ] [ counts>> ] tri [ clone ] tri@
 101     disjoint-set boa ;
 102
 103 : assoc>disjoint-set ( assoc -- disjoint-set )
 104     <disjoint-set>
 105     [ '[ drop _ add-atom ] assoc-each ]
 106     [ '[ _ equate ] assoc-each ]
 107     [ nip ]
 108     2tri ;