LeetCode/go/maximum-sum-of-3-non-overlapping-subarrays.go

package main

type Solution struct {
	Sum   int
	Index int
}

func (s *Solution) Update(previous *Solution, sum, index int) {
	if sum > previous.Sum {
		s.Sum = sum
		s.Index = index
		return
	}

	s.Sum = previous.Sum
	s.Index = previous.Index
}

func maxSumOfThreeSubarrays(nums []int, k int) []int {
	n := len(nums)

	prefix := make([]int, n+1)
	for i := 1; i <= n; i++ {
		prefix[i] = prefix[i-1] + nums[i-1]
	}

	best := make([][]Solution, 4)
	for i, _ := range best {
		best[i] = make([]Solution, n+1)
	}

	for subarrays := 1; subarrays <= 3; subarrays++ {
		for end := k * subarrays; end <= n; end++ {
			current := prefix[end] - prefix[end-k] + best[subarrays-1][end-k].Sum
			best[subarrays][end].Update(&best[subarrays][end-1], current, end-k)
		}
	}

	result := make([]int, 3)
	end := n
	for index := 3; index >= 1; index-- {
		result[index-1] = best[index][end].Index
		end = result[index-1]
	}

	return result
}