71 lines
1.4 KiB
Go
71 lines
1.4 KiB
Go
// Package q3666 implements a solution for https://leetcode.com/problems/minimum-operations-to-equalize-binary-string/
|
|
package q3666
|
|
|
|
import (
|
|
"strings"
|
|
|
|
rbt "github.com/emirpasic/gods/v2/trees/redblacktree"
|
|
)
|
|
|
|
type void struct{}
|
|
|
|
const INF = 1<<63 - 1
|
|
|
|
var (
|
|
qBuf = make([]int, 0, 1e5+1)
|
|
nOpsBuf = make([]int, 0, 1e5+1)
|
|
)
|
|
|
|
func minOperations(s string, k int) int {
|
|
zeros := strings.Count(s, "0")
|
|
length := len(s)
|
|
|
|
if zeros%2 == 1 && k%2 == 0 {
|
|
return -1
|
|
}
|
|
if k == length && zeros != 0 && zeros != length {
|
|
return -1
|
|
}
|
|
if zeros%k == 0 {
|
|
return zeros / k
|
|
}
|
|
|
|
nOps := nOpsBuf[:length+1]
|
|
sets := [2]*rbt.Tree[int, void]{
|
|
rbt.New[int, void](), rbt.New[int, void](),
|
|
}
|
|
for i := range length + 1 {
|
|
sets[i%2].Put(i, void{})
|
|
nOps[i] = INF
|
|
}
|
|
|
|
nOps[zeros] = 0
|
|
sets[zeros%2].Remove(zeros)
|
|
queue := qBuf[:0]
|
|
queue = append(queue, zeros)
|
|
|
|
for ; len(queue) > 0; queue = queue[1:] {
|
|
z := queue[0]
|
|
par := (z + k) % 2
|
|
|
|
minFlipZ := max(k-(length-z), 0)
|
|
maxFlipZ := min(k, z)
|
|
minNewZ := z + k - 2*maxFlipZ
|
|
maxNewZ := z + k - 2*minFlipZ
|
|
|
|
next, found := sets[par].Ceiling(minNewZ)
|
|
for ; found && next.Key <= maxNewZ; next, found = sets[par].Ceiling(minNewZ) {
|
|
nextZ := next.Key
|
|
sets[par].Remove(nextZ)
|
|
|
|
nOps[nextZ] = nOps[z] + 1
|
|
if nextZ%k == 0 {
|
|
return nOps[nextZ] + nextZ/k
|
|
}
|
|
queue = append(queue, nextZ)
|
|
}
|
|
}
|
|
return -1
|
|
}
|
|
|
|
var _ = minOperations
|