"Shortest Subarray with Sum at Least K" hard task No 862 accepted solution

862-shortest-subarray-with-sum-at-least-k.go

@@ -1,30 +1,97 @@
 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 := 999999
+	sz := math.MaxInt64
-	prefix := make([]int, len(nums)+1)
+	prefixSlice := make(PrefixHeap, len(nums)+1)
+	prefixHeap := PrefixHeap{}
 
-	for i := 1; i <= len(nums); i++ {
+	prefixSlice[0].sum = int64(nums[0])
-		prefix[i] = prefix[i-1] + nums[i-1]
+	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++ {
-		for j := i; j < len(nums); j++ {
+
-			sum := prefix[j+1] - prefix[i]
+		//fmt.Printf("---------------- %d ----------------\n", i)
-			if sum >= k {
+		//fmt.Println(prefixHeap)
-				l := j - i + 1
+		//if (0 != prefixHeap.Len()) && ((prefixSlice[i].sum - prefixHeap[0].sum) < int64(k)) {
-				if sz > l {
+		//	fmt.Printf("prefixHeap.Len() = %d, i = %d, prefixSlice[%d].sum - prefixHeap[0].sum = %d, k = %d\n",
-					sz = l
+		//		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
 			}
 		}
 
-	if sz == 999999 {
+		heap.Push(&prefixHeap, prefixSlice[i])
+		//fmt.Println(prefixHeap)
+		//fmt.Printf("---------------- %d ----------------\n", i)
+	}
+
+	if sz == math.MaxInt64 {
 		sz = -1
 	}