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

package main

import "sync"

func fromRightToLeft(fruityPrefix []int, startPos int, k int, leftLimit int, rightLimit int, pwg *sync.WaitGroup, presult *int) {
	calcFruitySliceSum := func(begin int, end int) int {
		begin--
		if begin < 0 {
			return fruityPrefix[end]
		}

		return fruityPrefix[end] - fruityPrefix[begin]
	}

	windowSize := k / 2
	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 + delta // for odd numbers of k we need to adjust wnd sz +1

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

	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
	}

	*presult = maxFruits
	pwg.Done()
}

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

		return fruityPrefix[end] - fruityPrefix[begin]
	}

	maxFruits := calcFruitySliceSum(leftLimit, startPos)

	windowSize := k / 2
	delta := k - windowSize*2
	windowEnd := startPos + delta
	windowBegin := windowEnd - windowSize - delta

	if windowBegin < leftLimit {
		windowEnd += 2 * (leftLimit - windowBegin)
		windowBegin = leftLimit
		windowSize = windowEnd - delta - 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
}

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)

	maxFruits := 0

	wg := new(sync.WaitGroup)
	wg.Add(1)

	go fromRightToLeft(fruityPrefix, startPos, k, leftLimit, rightLimit, wg, &maxFruits)
	alternativeMaxFruits := fromLeftToRight(fruityPrefix, startPos, k, leftLimit, rightLimit)

	wg.Wait()

	if maxFruits > alternativeMaxFruits {
		return maxFruits
	}

	return alternativeMaxFruits
}