wine.go

package main

import (
    "fmt"
    "math/bits"
    "sort"
)

// ceilLog2 returns the smallest d such that 2**d >= m.
func ceilLog2(m int) int {
    if m <= 1 {
        return 0
    }
    // bits.Len(x-1) gives ceil(log2(x)) for x>0
    return bits.Len(uint(m - 1))
}

// permutations generates all permutations of 0..n-1.
func permutations(n int) [][]int {
    nums := make([]int, n)
    for i := 0; i < n; i++ {
        nums[i] = i
    }
    var res [][]int
    var generate func(int)
    generate = func(k int) {
        if k == n {
            p := make([]int, n)
            copy(p, nums)
            res = append(res, p)
            return
        }
        for i := k; i < n; i++ {
            nums[k], nums[i] = nums[i], nums[k]
            generate(k + 1)
            nums[k], nums[i] = nums[i], nums[k]
        }
    }
    generate(0)
    return res
}

type query struct {
    box int
    bit int
}

// buildInstance constructs perms, queries and ans matrix.
func buildInstance(n int) (perms [][]int, queries []query, ans [][]uint8) {
    perms = permutations(n)
    maxBit := bits.Len(uint(n-1)) - 1
    for i := 0; i < n; i++ {
        for k := 0; k <= maxBit; k++ {
            queries = append(queries, query{box: i, bit: k})
        }
    }
    T := len(perms)
    ans = make([][]uint8, len(queries))
    for qi, q := range queries {
        row := make([]uint8, T)
        for t, p := range perms {
            row[t] = uint8((p[q.box] >> q.bit) & 1)
        }
        ans[qi] = row
    }
    return perms, queries, ans
}

// greedyUpperBound computes a greedy splitter depth upper bound.
func greedyUpperBound(ans [][]uint8, Tidx []int) int {
    memo := make(map[string]int)
    var keyBuf []byte

    var dfs func([]int) int
    dfs = func(state []int) int {
        if len(state) <= 1 {
            return 0
        }
        // key
        keyBuf = keyBuf[:0]
        for _, v := range state {
            keyBuf = append(keyBuf, byte(v>>8), byte(v))
        }
        key := string(keyBuf)
        if v, ok := memo[key]; ok {
            return v
        }
        bestMax := int(^uint(0) >> 1) // max int
        bestZ := make([]int, 0)
        bestO := make([]int, 0)
        for _, row := range ans {
            z := make([]int, 0, len(state))
            o := make([]int, 0, len(state))
            for _, t := range state {
                if row[t] == 0 {
                    z = append(z, t)
                } else {
                    o = append(o, t)
                }
            }
            if len(z) == 0 || len(o) == 0 {
                continue
            }
            m := len(z)
            if len(o) > m {
                m = len(o)
            }
            if m < bestMax {
                bestMax = m
                bestZ = z
                bestO = o
            }
        }
        if bestMax == int(^uint(0)>>1) {
            memo[key] = 1_000_000_000
            return memo[key]
        }
        res := 1 + max(dfs(bestZ), dfs(bestO))
        memo[key] = res
        return res
    }
    return dfs(Tidx)
}

// bestCandidates returns informative queries for this state ordered by local balance.
func bestCandidates(state []int, orderedQ []int, ans [][]uint8) []int {
    type cand struct{
        maxSide int
        minSide int
        qi int
    }
    cands := make([]cand, 0, len(orderedQ))
    for _, qi := range orderedQ {
        row := ans[qi]
        zc, oc := 0, 0
        for _, t := range state {
            if row[t] == 0 {
                zc++
            } else {
                oc++
            }
        }
        if zc == 0 || oc == 0 {
            continue
        }
        mx := zc
        if oc > mx { mx = oc }
        mn := zc
        if oc < mn { mn = oc }
        cands = append(cands, cand{maxSide: mx, minSide: -mn, qi: qi})
    }
    sort.Slice(cands, func(i, j int) bool {
        if cands[i].maxSide != cands[j].maxSide {
            return cands[i].maxSide < cands[j].maxSide
        }
        return cands[i].minSide < cands[j].minSide
    })
    res := make([]int, len(cands))
    for i, c := range cands {
        res[i] = c.qi
    }
    return res
}

// optimalDepthAndFirstQuery replicates the Python logic and returns the result map.
func optimalDepthAndFirstQuery(n int) map[string]any {
    perms, queries, ans := buildInstance(n)
    T := len(perms)
    Tidx := make([]int, T)
    for i := range Tidx { Tidx[i] = i }

    lb := ceilLog2(T)
    ub := greedyUpperBound(ans, Tidx)

    // order queries globally by balance on full set
    type score struct{
        maxSide int
        minSide int
        qi int
    }
    scores := make([]score, 0, len(ans))
    for qi, row := range ans {
        z := 0
        for _, t := range Tidx { if row[t] == 0 { z++ } }
        o := T - z
        if z == 0 || o == 0 { continue }
        mn := z
        if o < mn { mn = o }
        mx := z
        if o > mx { mx = o }
        scores = append(scores, score{maxSide: mx, minSide: -mn, qi: qi})
    }
    sort.Slice(scores, func(i, j int) bool {
        if scores[i].maxSide != scores[j].maxSide {
            return scores[i].maxSide < scores[j].maxSide
        }
        return scores[i].minSide < scores[j].minSide
    })
    orderedQ := make([]int, len(scores))
    for i, s := range scores { orderedQ[i] = s.qi }

    // feasibility with memoization
    type memoKey struct{
        // store state as sorted ints and depth
        d int
        key string
    }
    feasMemo := make(map[memoKey]bool)
    var keyBuf []byte
    makeKey := func(state []int) string {
        keyBuf = keyBuf[:0]
        for _, v := range state {
            keyBuf = append(keyBuf, byte(v>>8), byte(v))
        }
        return string(keyBuf)
    }

    var feasible func([]int, int) bool
    feasible = func(state []int, d int) bool {
        m := len(state)
        if m <= 1 { return true }
        if d < ceilLog2(m) { return false }
        k := memoKey{d: d, key: makeKey(state)}
        if v, ok := feasMemo[k]; ok { return v }
        for _, qi := range bestCandidates(state, orderedQ, ans) {
            row := ans[qi]
            zeros := make([]int, 0, m)
            ones := make([]int, 0, m)
            for _, t := range state {
                if row[t] == 0 { zeros = append(zeros, t) } else { ones = append(ones, t) }
            }
            if len(zeros) == 0 || len(ones) == 0 { continue }
            if feasible(zeros, d-1) && feasible(ones, d-1) {
                feasMemo[k] = true
                return true
            }
        }
        feasMemo[k] = false
        return false
    }

    var opt *int
    for d := lb; d <= ub; d++ {
        if feasible(Tidx, d) {
            val := d
            opt = &val
            break
        }
    }

    var firstQ *query
    if opt != nil {
        state := append([]int(nil), Tidx...)
        d := *opt
        for _, qi := range bestCandidates(state, orderedQ, ans) {
            row := ans[qi]
            var z, o []int
            for _, t := range state {
                if row[t] == 0 { z = append(z, t) } else { o = append(o, t) }
            }
            if len(z) > 0 && len(o) > 0 && feasible(z, d-1) && feasible(o, d-1) {
                q := queries[qi]
                firstQ = &q
                break
            }
        }
    }

    var fq any
    if firstQ != nil {
        fq = [2]int{firstQ.box, firstQ.bit}
    } else {
        fq = nil
    }

    return map[string]any{
        "n": n,
        "lower_bound": lb,
        "upper_bound_greedy": ub,
        "optimal_worst_case_depth": func() any { if opt==nil { return nil }; return *opt }(),
        "first_query": fq,
        "queries_indexed_as": "tuple(box, bit)",
    }
}

func max(a, b int) int { if a>b { return a }; return b }

func main() {
    res := optimalDepthAndFirstQuery(8)
    fmt.Println(res)
}