add new solutions

solutions/36/q3600/solution.go

@@ -0,0 +1,77 @@
+// Package q3600 implements a solution for https://leetcode.com/problems/maximize-spanning-tree-stability-with-upgrades/
+package q3600
+
+import (
+	"math"
+	"slices"
+)
+
+func unionFind(unions []int, a int) int {
+	if unions[a] == a {
+		return a
+	} else {
+		root := unionFind(unions, unions[a])
+		unions[a] = root
+		return root
+	}
+}
+
+func unionJoin(unions []int, a, b int) bool {
+	a, b = min(a, b), max(a, b) // suboptimal dirty impl. ideally, order by size of union
+	uA, uB := unionFind(unions, a), unionFind(unions, b)
+	if uA != uB {
+		unions[uB] = uA
+		return true
+	}
+	return false // cannot be joined. already in the same set.
+}
+
+func maxStability(n int, edges [][]int, k int) int {
+	unions := make([]int, n)
+	for i := range unions {
+		unions[i] = i
+	}
+
+	selected := make([]int, 0, len(edges)+1)
+	nSel, minStr := 0, math.MaxInt
+	for _, e := range edges {
+		if e[3] == 1 { // must include
+			nSel++
+			minStr = min(minStr, e[2])
+			if !unionJoin(unions, e[0], e[1]) {
+				return -1
+			}
+		}
+	}
+	slices.SortFunc(edges, func(a, b []int) int {
+		return b[2] - a[2] // sort by strength desc
+	})
+
+	for i := range edges {
+		if nSel >= n-1 {
+			break
+		}
+		if edges[i][3] == 1 { // must include, already done.
+			continue
+		}
+		uA, uB := unionFind(unions, edges[i][0]), unionFind(unions, edges[i][1])
+		if uA == uB { // already in the same tree
+			continue
+		}
+
+		nSel++
+		selected = append(selected, edges[i][2])
+		unionJoin(unions, uA, uB)
+	}
+
+	if nSel != n-1 {
+		return -1
+	}
+	for i := len(selected) - 1; i >= max(0, len(selected)-k); i-- {
+		selected[i] *= 2
+	}
+	selected = append(selected, minStr)
+	return slices.Min(selected)
+}
+
+var _ = maxStability