LeetCode 带限制的子序列和

/**
LeetCode 带限制的子序列和
由题意可知，第i位的最大子序和，是在区间[i-k, i-1]中的最大值加上nums[i]，但是若得出的结果，比nums[i]小，便舍弃。
转移方程：
dp[i] = Max(nums[i], dp[i-k...i-1] + nums[i])
初始化：
dp[0] = nums[0]
利用堆heap维护dp[i-k...i-1]的递减序列
*/
func constrainedSubsetSum(nums []int, k int) int {
	numLen := len(nums)
	dp := make([]int, numLen+1)
	dp[0] = nums[0]
	heap := heap{arr: make([]node, numLen+1), sz:0}
	heap.push(0, dp[0])
	ans := dp[0]
	for i:=1; i<numLen; i++ {
		for heap.top().p < i-k && heap.size() != 0{
			heap.pop()
		}
		dp[i] = Max(nums[i], heap.top().v + nums[i])
		heap.push(i, dp[i])
		ans = Max(ans, dp[i])
	}
	return ans
}

func Max(i, j int) int {
	if i > j {
		return i
	}
	return j
}

func Min(i, j int) int {
	if i > j {
		return j
	}
	return i
}

type node struct {
	v, p int
}
type heap struct {
	arr []node
	sz int
}

func (h *heap) push(p, v int) {
	h.arr[h.sz].p = p
	h.arr[h.sz].v = v
	h.sz++
	h.up()
}

func (h *heap) clear() {
	h.sz = 0
}

func (h *heap) top() node {
	return h.arr[0]
}

func (h *heap) pop() {
	if h.sz > 0 {
		h.arr[0], h.arr[h.sz-1] = h.arr[h.sz-1], h.arr[0]
		h.sz--
		h.down()
	}
	return
}

func (h *heap) size() int {
	return h.sz
}

func (h *heap) down() {
	p := 0
	for p < h.sz {
		l := p*2 + 1
		r := p*2 + 2
		if l >= h.sz && r >= h.sz {
			break
		}
		pi := l
		if l < h.sz && r < h.sz {
			if h.arr[l].v > h.arr[r].v {
				pi = l
			} else {
				pi = r
			}
		}

		if h.arr[pi].v > h.arr[p].v {
			h.arr[pi], h.arr[p] = h.arr[p], h.arr[pi]
			p = pi
		} else {
			break
		}
	}
	return
}

func (h *heap) up() {
	j := h.sz-1
	for {
		i := (j-1) / 2
		if i == j || h.arr[j].v < h.arr[i].v {
			break
		}
		h.arr[i], h.arr[j] = h.arr[j], h.arr[i]
		j = i
	}
}