diff --git a/algo/unionfind.go b/algo/unionfind.go new file mode 100644 index 0000000..6fb5359 --- /dev/null +++ b/algo/unionfind.go @@ -0,0 +1,45 @@ +package algo + +type Subset struct { + parent string + rank int +} + +// find is a recursive function that finds the root of the subset that element 'e' belongs to. +// It uses path compression technique to optimize future find operations. +// The subsets parameter is a map that stores the subsets with their parent information. +// The e parameter is the element for which the root is to be found. +// The function returns the root of the subset that 'e' belongs to. +// see https://ondrej-kvasnovsky-2.gitbook.io/algorithms/finding/union-find-algorithms +// or https://jojozhuang.github.io/algorithm/algorithm-union-find/ +func find(subsets map[string]*Subset, e string) string { + // find root and make root as parent of e + // (path compression) + if subsets[e].parent != e { + subsets[e].parent = find(subsets, subsets[e].parent) + } + return subsets[e].parent +} + +// Union merges two subsets in the union-find data structure. +// It takes a map of subsets, and the names of the two subsets to be merged. +// The function finds the root of each subset and performs the union operation based on the rank of the subsets. +func Union(subsets map[string]*Subset, x, y string) { + xroot := find(subsets, x) + yroot := find(subsets, y) + + // Attach smaller rank tree under root of high + // rank tree (Union by Rank) + if subsets[xroot].rank < subsets[yroot].rank { + subsets[xroot].parent = yroot + } else { + if subsets[xroot].rank > subsets[yroot].rank { + subsets[yroot].parent = xroot + } else { + // If ranks are same, then make one as root and + // increment its rank by one + subsets[yroot].parent = xroot + subsets[xroot].rank++ + } + } +} diff --git a/algo/unionfind_test.go b/algo/unionfind_test.go new file mode 100644 index 0000000..6c26442 --- /dev/null +++ b/algo/unionfind_test.go @@ -0,0 +1,111 @@ +package algo + +import ( + "testing" +) + +func TestFind(t *testing.T) { + subsets := map[string]*Subset{ + "a": {parent: "a"}, + "b": {parent: "a"}, + "c": {parent: "b"}, + "d": {parent: "c"}, + } + + tests := []struct { + name string + subsets map[string]*Subset + element string + expected string + }{ + { + name: "element is root", + subsets: subsets, + element: "a", + expected: "a", + }, + { + name: "element has parent", + subsets: subsets, + element: "b", + expected: "a", + }, + { + name: "element has grandparent", + subsets: subsets, + element: "c", + expected: "a", + }, + { + name: "element has great-grandparent", + subsets: subsets, + element: "d", + expected: "a", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := find(tt.subsets, tt.element) + if got != tt.expected { + t.Errorf("find() = %v, want %v", got, tt.expected) + } + }) + } +} + +func TestUnion(t *testing.T) { + subsets := map[string]*Subset{ + "a": {parent: "a"}, + "b": {parent: "a"}, + "c": {parent: "b"}, + "d": {parent: "c"}, + } + + tests := []struct { + name string + subsets map[string]*Subset + x string + y string + expectedXRoot string + }{ + { + name: "x and y are already in the same subset", + subsets: subsets, + x: "a", + y: "b", + expectedXRoot: "a", + }, + { + name: "x and y are in different subsets", + subsets: subsets, + x: "b", + y: "c", + expectedXRoot: "a", + }, + { + name: "x and y are in different subsets with different ranks", + subsets: subsets, + x: "c", + y: "d", + expectedXRoot: "a", + }, + { + name: "x and y are in different subsets with same rank", + subsets: subsets, + x: "a", + y: "d", + expectedXRoot: "a", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + Union(tt.subsets, tt.x, tt.y) + got := find(tt.subsets, tt.x) + if got != tt.expectedXRoot { + t.Errorf("Union() failed, got x root = %v, want %v", got, tt.expectedXRoot) + } + }) + } +}