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 }