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 ;
   5
   6 IN: disjoint-sets
   7
   8 TUPLE: disjoint-set
   9 { parents hashtable read-only }
  10 { ranks hashtable read-only }
  11 { counts hashtable read-only } ;
  12
  13 <PRIVATE
  14
  15 : count ( a disjoint-set -- n )
  16     counts>> at ; inline
  17
  18 : add-count ( p a disjoint-set -- )
  19     [ count [ + ] curry ] keep counts>> swap change-at ; inline
  20
  21 : parent ( a disjoint-set -- p )
  22     parents>> at ; inline
  23
  24 : set-parent ( p a disjoint-set -- )
  25     parents>> set-at ; inline
  26
  27 : link-sets ( p a disjoint-set -- )
  28     [ set-parent ] [ add-count ] 3bi ; inline
  29
  30 : rank ( a disjoint-set -- r )
  31     ranks>> at ; inline
  32
  33 : inc-rank ( a disjoint-set -- )
  34     ranks>> [ 1+ ] change-at ; inline
  35
  36 : representative? ( a disjoint-set -- ? )
  37     dupd parent = ; inline
  38
  39 PRIVATE>
  40
  41 GENERIC: representative ( a disjoint-set -- p )
  42
  43 M: disjoint-set representative
  44     2dup representative? [ drop ] [
  45         [ [ parent ] keep representative dup ] 2keep set-parent
  46     ] if ;
  47
  48 <PRIVATE
  49
  50 : representatives ( a b disjoint-set -- r r )
  51     [ representative ] curry bi@ ; inline
  52
  53 : ranks ( a b disjoint-set -- r r )
  54     [ rank ] curry bi@ ; inline
  55
  56 :: branch ( a b neg zero pos -- )
  57     a b = zero [ a b < neg pos if ] if ; inline
  58
  59 PRIVATE>
  60
  61 : <disjoint-set> ( -- disjoint-set )
  62     H{ } clone H{ } clone H{ } clone disjoint-set boa ;
  63
  64 GENERIC: add-atom ( a disjoint-set -- )
  65
  66 M: disjoint-set add-atom
  67     [ dupd parents>> set-at ]
  68     [ 0 -rot ranks>> set-at ]
  69     [ 1 -rot counts>> set-at ]
  70     2tri ;
  71
  72 GENERIC: equiv-set-size ( a disjoint-set -- n )
  73
  74 M: disjoint-set equiv-set-size [ representative ] keep count ;
  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 ;