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