package main func MaxTotalFruits(fruits [][]int, startPos int, k int) int { if len(fruits) == 0 || k < 0 { return 0 } 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) windowBegin := startPos - (k - windowSize*2) // 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 { windowEnd = rightLimit } 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 } 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 { windowBegin = leftLimit } 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 } return maxFruits }