100 lines
2.1 KiB
Go
100 lines
2.1 KiB
Go
package main
|
|
|
|
import (
|
|
"container/heap"
|
|
"math"
|
|
)
|
|
|
|
type Prefix struct {
|
|
sum int64
|
|
index int
|
|
}
|
|
|
|
type PrefixHeap []Prefix
|
|
|
|
func (p *PrefixHeap) Len() int {
|
|
return len(*p)
|
|
}
|
|
|
|
func (p *PrefixHeap) Swap(i, j int) {
|
|
(*p)[i], (*p)[j] = (*p)[j], (*p)[i]
|
|
}
|
|
|
|
func (p *PrefixHeap) Less(i, j int) bool {
|
|
if (*p)[i].sum != (*p)[j].sum {
|
|
return (*p)[i].sum < (*p)[j].sum
|
|
}
|
|
return (*p)[i].index < (*p)[j].index
|
|
}
|
|
|
|
func (p *PrefixHeap) Push(x interface{}) {
|
|
it := x.(Prefix)
|
|
*p = append(*p, it)
|
|
}
|
|
|
|
func (p *PrefixHeap) Pop() interface{} {
|
|
old := *p
|
|
n := len(old)
|
|
item := old[n-1]
|
|
*p = old[0 : n-1]
|
|
return item
|
|
}
|
|
|
|
func ShortestSubarray(nums []int, k int) int {
|
|
if len(nums) == 0 {
|
|
return -1
|
|
}
|
|
|
|
sz := math.MaxInt64
|
|
prefixSlice := make(PrefixHeap, len(nums)+1)
|
|
prefixHeap := PrefixHeap{}
|
|
|
|
prefixSlice[0].sum = int64(nums[0])
|
|
prefixSlice[0].index = 0
|
|
|
|
for i := 1; i < len(nums); i++ {
|
|
prefixSlice[i].sum = prefixSlice[i-1].sum + int64(nums[i])
|
|
prefixSlice[i].index = i
|
|
}
|
|
|
|
//fmt.Print("Initial state. prefixSlice = ")
|
|
//fmt.Println(prefixSlice)
|
|
|
|
heap.Push(&prefixHeap, Prefix{sum: 0, index: -1})
|
|
|
|
for i := 0; i < len(nums); i++ {
|
|
|
|
//fmt.Printf("---------------- %d ----------------\n", i)
|
|
//fmt.Println(prefixHeap)
|
|
//if (0 != prefixHeap.Len()) && ((prefixSlice[i].sum - prefixHeap[0].sum) < int64(k)) {
|
|
// fmt.Printf("prefixHeap.Len() = %d, i = %d, prefixSlice[%d].sum - prefixHeap[0].sum = %d, k = %d\n",
|
|
// prefixHeap.Len(), i, i, prefixSlice[i].sum-prefixHeap[0].sum, k)
|
|
//}
|
|
|
|
for (0 != prefixHeap.Len()) && ((prefixSlice[i].sum - prefixHeap[0].sum) >= int64(k)) {
|
|
// fmt.Printf("prefixHeap.Len() = %d, i = %d, prefixSlice[%d].sum - prefixHeap[0].sum = %d, k = %d\n",
|
|
// prefixHeap.Len(), i, i, prefixSlice[i].sum-prefixHeap[0].sum, k)
|
|
|
|
p := heap.Pop(&prefixHeap).(Prefix)
|
|
ln := i - p.index
|
|
|
|
//fmt.Println("p = ", p, ", ln = ", ln, ", sz = ", sz)
|
|
|
|
if ln < sz {
|
|
//fmt.Println("==== (ln < sz) ====>> (sz = ln) ====")
|
|
sz = ln
|
|
}
|
|
}
|
|
|
|
heap.Push(&prefixHeap, prefixSlice[i])
|
|
//fmt.Println(prefixHeap)
|
|
//fmt.Printf("---------------- %d ----------------\n", i)
|
|
}
|
|
|
|
if sz == math.MaxInt64 {
|
|
sz = -1
|
|
}
|
|
|
|
return sz
|
|
}
|