diff --git a/go/remove-max-number-of-edges-to-keep-graph-fully-traversable.go b/go/remove-max-number-of-edges-to-keep-graph-fully-traversable.go new file mode 100644 index 0000000..174887a --- /dev/null +++ b/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 +}