go: add «1579. Remove Max Number of Edges to Keep Graph Fully Traversable»

Signed-off-by: Matej Focko <me@mfocko.xyz>

go/remove-max-number-of-edges-to-keep-graph-fully-traversable.go

@@ -0,0 +1,104 @@
+package main
+
+import (
+	"cmp"
+	"slices"
+)
+
+type UnionFind struct {
+	parent []int
+	rank   []int
+	count  int
+}
+
+func (uf *UnionFind) Count() int {
+	return uf.count
+}
+
+func (uf *UnionFind) find(u int) int {
+	if uf.parent[u] != u {
+		uf.parent[u] = uf.find(uf.parent[u])
+	}
+
+	return uf.parent[u]
+}
+
+func (uf *UnionFind) Union(u, v int) bool {
+	rootU := uf.find(u)
+	rootV := uf.find(v)
+
+	if rootU == rootV {
+		return false
+	}
+
+	if uf.rank[rootU] > uf.rank[rootV] {
+		uf.parent[rootV] = rootU
+	} else if uf.rank[rootU] < uf.rank[rootV] {
+		uf.parent[rootU] = rootV
+	} else {
+		uf.parent[rootV] = rootU
+		uf.rank[rootU]++
+	}
+
+	uf.count--
+	return true
+}
+
+func makeUnionFind(n int) UnionFind {
+	parent := make([]int, n)
+	for i := range n {
+		parent[i] = i
+	}
+
+	return UnionFind{
+		parent: parent,
+		rank:   make([]int, n),
+		count:  n,
+	}
+}
+
+const (
+	Alice int = 1
+	Bob       = 2
+	Both      = 3
+)
+
+// by edge type — preferring «Both» type
+func byEdgeType(l, r []int) int {
+	return -cmp.Compare(l[0], r[0])
+}
+
+func maxNumEdgesToRemove(n int, edges [][]int) int {
+	slices.SortFunc(edges, byEdgeType)
+
+	alice := makeUnionFind(n)
+	bob := makeUnionFind(n)
+
+	added := 0
+	for _, edge := range edges {
+		t, u, v := edge[0], edge[1]-1, edge[2]-1
+
+		switch t {
+		case Alice:
+			if alice.Union(u, v) {
+				added++
+			}
+		case Bob:
+			if bob.Union(u, v) {
+				added++
+			}
+		case Both:
+			connectsAlice := alice.Union(u, v)
+			connectsBob := bob.Union(u, v)
+			if connectsAlice || connectsBob {
+				added++
+			}
+		}
+	}
+
+	if alice.Count() > 1 || bob.Count() > 1 {
+		return -1
+	}
+
+	return len(edges) - added
+}