// 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