leetcode-go/leetcode/hard/689.maximum-sum-of-3-non-overlapping-subarrays.go

func maxSumOfThreeSubarrays(nums []int, k int) []int {
	n := len(nums)
	sum := make([]int, n+1)
	for i := 1; i <= n; i++ {
		sum[i] = nums[i-1] + sum[i-1]
	}
	dp := make([][][4]int, 3)
	// dp[i][j] means the maximum sum of i+1 sub arrays until j,
	// dp[i][j] = [4]int{max, i1, i2, i3}
	//                        '-> corresponding index of the maximum
	for i := range dp {
		dp[i] = make([][4]int, n)
	}
	for i := 0; i < 3; i++ {
		jmax := n - (3-i)*k
		for j := i * k; j <= jmax; j++ {
			s := sum[j+k] - sum[j]
			if i == 0 {
				if 0 < j && s <= dp[i][j-1][0] {
					dp[i][j] = dp[i][j-1]
				} else {
					dp[i][j] = [4]int{s, j, 0, 0}
				}
				continue
			}
			if ns := dp[i-1][j-k][0] + s; ns <= dp[i][j-1][0] {
				dp[i][j] = dp[i][j-1]
			} else {
				dp[i][j] = dp[i-1][j-k]
				dp[i][j][0], dp[i][j][i+1] = ns, j
			}
		}
	}
	return dp[2][n-k][1:]
}