leetcode/2106-max-fruits-harvested-in-k-steps.go

package main

func MaxTotalFruits(fruits [][]int, startPos int, k int) int {
	if len(fruits) == 0 || k < 0 {
		return 0
	}

	//fmt.Println("startPos = ", startPos)
	//fmt.Println("steps(k) = ", k)

	totalLength := fruits[len(fruits)-1][0] + 1
	//fmt.Println("totalLength = ", totalLength)

	if startPos >= totalLength {
		totalLength = startPos + 1
	}

	fruitySlice := make([]int, totalLength)
	fruityPrefix := make([]int, totalLength)

	for _, fruit := range fruits {
		fruitySlice[fruit[0]] = fruit[1]
	}
	//fmt.Println("fruitySlice = ", fruitySlice)

	if k == 0 {
		return fruitySlice[startPos]
	}

	prefixSum := 0
	for i, fruit := range fruitySlice {
		prefixSum += fruit
		fruityPrefix[i] = prefixSum
	}
	//fmt.Println("fruityPrefix = ", fruityPrefix)

	leftLimit := startPos - k
	if leftLimit < 0 {
		leftLimit = 0
	}
	//fmt.Println("leftLimit = ", leftLimit)

	rightLimit := startPos + k
	if rightLimit >= totalLength {
		rightLimit = totalLength - 1
	}
	//fmt.Println("rightLimit = ", rightLimit)

	windowSize := k / 2

	calcFruitySliceSum := func(begin int, end int) int {
		begin--
		if begin < 0 {
			return fruityPrefix[end]
		}

		return fruityPrefix[end] - fruityPrefix[begin]
	}

	maxFruits := calcFruitySliceSum(startPos, rightLimit)

	delta := k - windowSize*2
	windowBegin :=
		startPos - delta // in case number of steps cannot be fully divided by 2, start position is shifted 1 position to the left
	windowEnd :=
		windowBegin + windowSize + (k - windowSize*2) // for odd numbers of k we need to adjust wnd sz +1

	if windowEnd > rightLimit {
		windowBegin -= 2 * (windowEnd - rightLimit)
		windowEnd = rightLimit
		windowSize = windowEnd - windowBegin - (k - windowSize*2)
	}

	endCycle := false
	for !endCycle {
		if windowBegin <= leftLimit {
			windowBegin = leftLimit
			endCycle = true
		}

		localMaxFruits := calcFruitySliceSum(windowBegin, windowEnd)
		if localMaxFruits > maxFruits {
			maxFruits = localMaxFruits
		}

		windowSize++
		windowEnd--
		windowBegin = windowEnd - windowSize - delta
	}

	alternativeMaxFruits := calcFruitySliceSum(leftLimit, startPos)
	if alternativeMaxFruits > maxFruits {
		maxFruits = alternativeMaxFruits
	}

	windowSize = k / 2
	windowEnd = startPos + (k - windowSize*2)
	windowBegin = windowEnd - windowSize - (k - windowSize*2)

	if windowBegin < leftLimit {
		windowEnd += 2 * (leftLimit - windowBegin)
		windowBegin = leftLimit
		windowSize = windowEnd - (k - windowSize*2) - windowBegin
	}

	endCycle = false
	for !endCycle {
		if windowEnd >= rightLimit {
			windowEnd = rightLimit
			endCycle = true
		}

		localMaxFruits := calcFruitySliceSum(windowBegin, windowEnd)
		if localMaxFruits > maxFruits {
			maxFruits = localMaxFruits
		}

		windowSize++
		windowBegin++
		windowEnd = windowBegin + windowSize + delta
	}

	return maxFruits
}