77 lines
1.6 KiB
Go
77 lines
1.6 KiB
Go
// 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
|