|
| 1 | + |
| 2 | +import option::none; |
| 3 | +import option::some; |
| 4 | + |
| 5 | + |
| 6 | +// A very naive implementation of union-find with unsigned integer nodes. |
| 7 | +// Maintains the invariant that the root of a node is always equal to or less |
| 8 | +// than the node itself. |
| 9 | +type node = option::t[uint]; |
| 10 | + |
| 11 | +type ufind = rec(mutable node[mutable] nodes); |
| 12 | + |
| 13 | +fn make() -> ufind { |
| 14 | + ret rec(mutable nodes=~[mutable]); |
| 15 | +} |
| 16 | + |
| 17 | +fn make_set(&ufind ufnd) -> uint { |
| 18 | + auto idx = ivec::len(ufnd.nodes); |
| 19 | + ufnd.nodes += ~[mutable none[uint]]; |
| 20 | + ret idx; |
| 21 | +} |
| 22 | + |
| 23 | + |
| 24 | +/// Creates sets as necessary to ensure that least `n` sets are present in the |
| 25 | +/// data structure. |
| 26 | +fn grow(&ufind ufnd, uint n) { |
| 27 | + while (set_count(ufnd) < n) { make_set(ufnd); } |
| 28 | +} |
| 29 | + |
| 30 | +fn find(&ufind ufnd, uint n) -> uint { |
| 31 | + alt (ufnd.nodes.(n)) { |
| 32 | + case (none) { ret n; } |
| 33 | + case (some(?m)) { auto m_ = m; be find(ufnd, m_); } |
| 34 | + } |
| 35 | +} |
| 36 | + |
| 37 | +fn union(&ufind ufnd, uint m, uint n) { |
| 38 | + auto m_root = find(ufnd, m); |
| 39 | + auto n_root = find(ufnd, n); |
| 40 | + if (m_root < n_root) { |
| 41 | + ufnd.nodes.(n_root) = some[uint](m_root); |
| 42 | + } else if (m_root > n_root) { ufnd.nodes.(m_root) = some[uint](n_root); } |
| 43 | +} |
| 44 | + |
| 45 | +fn set_count(&ufind ufnd) -> uint { ret ivec::len[node](ufnd.nodes); } |
| 46 | + |
| 47 | + |
| 48 | +// Removes all sets with IDs greater than or equal to the given value. |
| 49 | +fn prune(&ufind ufnd, uint n) { |
| 50 | + // TODO: Use "slice" once we get rid of "mutable?" |
| 51 | + |
| 52 | + auto len = ivec::len[node](ufnd.nodes); |
| 53 | + while (len != n) { ivec::pop[node](ufnd.nodes); len -= 1u; } |
| 54 | +} |
0 commit comments