code wiki / (root) / sketch_union_find.nx

sketch_union_find.nx source

↩ module page · 152 lines · 4802 B

1// sketch_union_find.nx -- Disjoint Set Union (Union-Find) primitive. 2// 3// Foundational data structure for connectivity / clustering. Maintains 4// a partition of {0, 1, ..., n-1} into disjoint sets. Operations: 5// find(x): which set is x in? (returns representative element) 6// union(x, y): merge x's set with y's set 7// connected(x, y): same set? 8// size_of(x): cardinality of x's set 9// 10// OPTIMIZATIONS: 11// - Path compression: find(x) flattens the parent chain so subsequent 12// finds are O(1) on the same nodes. 13// - Union by size: smaller tree attaches under larger; keeps tree 14// height O(log n). 15// 16// With both optimizations: nearly-O(1) amortized per operation 17// (technically O(α(n)) inverse-Ackermann; practically constant). 18// 19// USE CASES: 20// - Kruskal's MST (connectivity tracking) 21// - online clustering (merge clusters as new edges arrive) 22// - dynamic connectivity (undirected only) 23// - compilation: equivalence classes for SSA, alias analysis 24// - image processing: connected components 25// 26// LOSSLESS-LANGUAGE DISCIPLINE: 27// Production tier, exact semantics. No envelope returned -- 28// operations are deterministic with O(α(n)) guarantee. 29 30import "syscalls.nx" 31import "sketch_types.nx" 32 33const NX_UF_MIN_N: i64 = 1 34const NX_UF_MAX_N: i64 = 100000000 // 100M 35 36struct UnionFind { 37 parent: *i64, // parent[i] is the parent of i (-1 for root) 38 size: *i64, // size of set rooted at i (only valid at roots) 39 n: i64, 40 n_sets: i64, // count of distinct sets 41} 42 43// === construction ================================================= 44// 45// Initially each element is in its own set: parent[i] = i, size[i] = 1. 46 47func nx_uf_alloc(n: i64) -> *UnionFind { 48 if n < NX_UF_MIN_N { return 0 as *UnionFind } 49 if n > NX_UF_MAX_N { return 0 as *UnionFind } 50 let raw: *u8 = sys_mmap(40) 51 let uf: *UnionFind = raw as *UnionFind 52 uf.parent = sys_mmap(n * 8) as *i64 53 uf.size = sys_mmap(n * 8) as *i64 54 var i: i64 = 0 55 while i < n { 56 uf.parent[i] = i 57 uf.size[i] = 1 58 i = i + 1 59 } 60 uf.n = n 61 uf.n_sets = n 62 return uf 63} 64 65// === find (with path compression) ================================ 66// 67// Walk parent chain to root. On the way back down, re-point every 68// node to the root for amortized O(1) future finds. 69 70func nx_uf_find(uf: *UnionFind, x: i64) -> i64 { 71 if x < 0 { return -1 } 72 if x >= uf.n { return -1 } 73 // Find root. 74 var root: i64 = x 75 var done: i64 = 0 76 while done == 0 { 77 let p: i64 = uf.parent[root] 78 if p == root { done = 1 } 79 if done == 0 { root = p } 80 } 81 // Path compression: re-point intermediate nodes to root. 82 var cur: i64 = x 83 while cur != root { 84 let next: i64 = uf.parent[cur] 85 uf.parent[cur] = root 86 cur = next 87 } 88 return root 89} 90 91// === union (by size) ============================================= 92// 93// Find roots of x and y; if different, attach smaller-size root under 94// larger. Update size of new root. Decrement n_sets. 95 96func nx_uf_union(uf: *UnionFind, x: i64, y: i64) -> i64 { 97 let rx: i64 = nx_uf_find(uf, x) 98 let ry: i64 = nx_uf_find(uf, y) 99 if rx < 0 { return -1 } 100 if ry < 0 { return -1 } 101 if rx == ry { return 0 } // already in same set; no-op 102 let sx: i64 = uf.size[rx] 103 let sy: i64 = uf.size[ry] 104 if sx < sy { 105 uf.parent[rx] = ry 106 uf.size[ry] = sx + sy 107 } 108 if sx >= sy { 109 uf.parent[ry] = rx 110 uf.size[rx] = sx + sy 111 } 112 uf.n_sets = uf.n_sets - 1 113 return 0 114} 115 116// === connected (predicate) ======================================= 117 118func nx_uf_connected(uf: *UnionFind, x: i64, y: i64) -> i64 { 119 let rx: i64 = nx_uf_find(uf, x) 120 let ry: i64 = nx_uf_find(uf, y) 121 if rx < 0 { return 0 } 122 if ry < 0 { return 0 } 123 if rx == ry { return 1 } 124 return 0 125} 126 127// === size-of-set ================================================= 128 129func nx_uf_size_of(uf: *UnionFind, x: i64) -> i64 { 130 let rx: i64 = nx_uf_find(uf, x) 131 if rx < 0 { return 0 } 132 return uf.size[rx] 133} 134 135func nx_uf_n_sets(uf: *UnionFind) -> i64 { 136 return uf.n_sets 137} 138 139// === query (typed envelope wrapper for "size of set containing x") = 140// 141// Exact semantics -- envelope is degenerate (zero error). 142 143func nx_uf_query_size(uf: *UnionFind, x: i64) -> *ApproxI64 { 144 let s: i64 = nx_uf_size_of(uf, x) 145 return nx_approx_new(s, NX_ENV_ABS, 0, 1000000000, 146 NX_MATURITY_PRODUCTION, 147 NX_ADV_HONEST) 148} 149 150func nx_uf_memory_bytes(uf: *UnionFind) -> i64 { 151 return 40 + uf.n * 16 152}