diff --git a/2106-max-fruits-harvested-in-k-steps.go b/2106-max-fruits-harvested-in-k-steps.go new file mode 100644 index 0000000..62bc885 --- /dev/null +++ b/2106-max-fruits-harvested-in-k-steps.go @@ -0,0 +1,114 @@ +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 +} diff --git a/2106-max-fruits-harvested-in-k-steps_test.go b/2106-max-fruits-harvested-in-k-steps_test.go new file mode 100644 index 0000000..564d57f --- /dev/null +++ b/2106-max-fruits-harvested-in-k-steps_test.go @@ -0,0 +1,16 @@ +package main + +import ( + "github.com/stretchr/testify/assert" + "testing" +) + +func TestMaxTotalFruits(t *testing.T) { + + /*assert.Equal(t, 9, MaxTotalFruits([][]int{{2, 8}, {6, 3}, {8, 6}}, 5, 4)) + assert.Equal(t, 14, MaxTotalFruits([][]int{{0, 9}, {4, 1}, {5, 7}, {6, 2}, {7, 4}, {10, 9}}, 5, 4)) + assert.Equal(t, 0, MaxTotalFruits([][]int{{0, 3}, {6, 4}, {8, 5}}, 3, 2)) + assert.Equal(t, 10000, MaxTotalFruits([][]int{{0, 10000}}, 200000, 200000)) + assert.Equal(t, 10000, MaxTotalFruits([][]int{{200000, 10000}}, 0, 200000))*/ + assert.Equal(t, 22, MaxTotalFruits([][]int{{1, 9}, {2, 10}, {3, 1}, {5, 6}, {6, 3}, {8, 2}, {9, 2}, {11, 4}, {18, 10}, {22, 8}, {25, 2}, {26, 2}, {30, 4}, {31, 5}, {33, 9}, {34, 1}, {39, 10}}, 19, 9)) +}