"Maximum Fruits Harvested After at Most K Steps" accepted solution

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

@@ -5,6 +5,9 @@ func MaxTotalFruits(fruits [][]int, startPos int, k int) int {
 		return 0
 	}
 
+	//fmt.Println("startPos = ", startPos)
+	//fmt.Println("steps(k) = ", k)
+
 	totalLength := fruits[len(fruits)-1][0] + 1
 	//fmt.Println("totalLength = ", totalLength)
 
@@ -56,11 +59,16 @@ func MaxTotalFruits(fruits [][]int, startPos int, k int) int {
 
 	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
+	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 {
+	if windowEnd > rightLimit {
+		windowBegin -= 2 * (windowEnd - rightLimit)
 		windowEnd = rightLimit
+		windowSize = windowEnd - windowBegin - (k - windowSize*2)
 	}
 
 	endCycle := false
@@ -77,7 +85,7 @@ func MaxTotalFruits(fruits [][]int, startPos int, k int) int {
 
 		windowSize++
 		windowEnd--
-		windowBegin = windowEnd - windowSize
+		windowBegin = windowEnd - windowSize - delta
 	}
 
 	alternativeMaxFruits := calcFruitySliceSum(leftLimit, startPos)
@@ -89,8 +97,10 @@ func MaxTotalFruits(fruits [][]int, startPos int, k int) int {
 	windowEnd = startPos + (k - windowSize*2)
 	windowBegin = windowEnd - windowSize - (k - windowSize*2)
 
-	if windowBegin <= leftLimit {
+	if windowBegin < leftLimit {
+		windowEnd += 2 * (leftLimit - windowBegin)
 		windowBegin = leftLimit
+		windowSize = windowEnd - (k - windowSize*2) - windowBegin
 	}
 
 	endCycle = false
@@ -107,7 +117,7 @@ func MaxTotalFruits(fruits [][]int, startPos int, k int) int {
 
 		windowSize++
 		windowBegin++
-		windowEnd = windowBegin + windowSize
+		windowEnd = windowBegin + windowSize + delta
 	}
 
 	return maxFruits