Skip to content

2106. Maximum Fruits Harvested After at Most K Steps 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
 public:
  int maxTotalFruits(vector<vector<int>>& fruits, int startPos, int k) {
    const int maxRight = max(startPos, fruits.back()[0]);
    int ans = 0;
    vector<int> amounts(1 + maxRight);
    vector<int> prefix(2 + maxRight);

    for (const vector<int>& f : fruits)
      amounts[f[0]] = f[1];

    partial_sum(amounts.begin(), amounts.end(), prefix.begin() + 1);

    auto getFruits = [&](int leftSteps, int rightSteps) {
      const int l = max(0, startPos - leftSteps);
      const int r = min(maxRight, startPos + rightSteps);
      return prefix[r + 1] - prefix[l];
    };

    // Go right first.
    const int maxRightSteps = min(maxRight - startPos, k);
    for (int rightSteps = 0; rightSteps <= maxRightSteps; ++rightSteps) {
      const int leftSteps = max(0, k - 2 * rightSteps);  // Turn left
      ans = max(ans, getFruits(leftSteps, rightSteps));
    }

    // Go left first.
    const int maxLeftSteps = min(startPos, k);
    for (int leftSteps = 0; leftSteps <= maxLeftSteps; ++leftSteps) {
      const int rightSteps = max(0, k - 2 * leftSteps);  // Turn right
      ans = max(ans, getFruits(leftSteps, rightSteps));
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
  public int maxTotalFruits(int[][] fruits, int startPos, int k) {
    final int maxRight = Math.max(startPos, fruits[fruits.length - 1][0]);
    int ans = 0;
    int[] amounts = new int[1 + maxRight];
    int[] prefix = new int[2 + maxRight];

    for (int[] f : fruits)
      amounts[f[0]] = f[1];

    for (int i = 0; i + 1 < prefix.length; ++i)
      prefix[i + 1] = prefix[i] + amounts[i];

    // Go right first.
    final int maxRightSteps = Math.min(maxRight - startPos, k);
    for (int rightSteps = 0; rightSteps <= maxRightSteps; ++rightSteps) {
      final int leftSteps = Math.max(0, k - 2 * rightSteps); // Turn left
      ans = Math.max(ans, getFruits(startPos, maxRight, leftSteps, rightSteps, prefix));
    }

    // Go left first.
    final int maxLeftSteps = Math.min(startPos, k);
    for (int leftSteps = 0; leftSteps <= maxLeftSteps; ++leftSteps) {
      final int rightSteps = Math.max(0, k - 2 * leftSteps); // Turn right
      ans = Math.max(ans, getFruits(startPos, maxRight, leftSteps, rightSteps, prefix));
    }

    return ans;
  }

  private int getFruits(int startPos, int maxRight, int leftSteps, int rightSteps, int[] prefix) {
    final int l = Math.max(0, startPos - leftSteps);
    final int r = Math.min(maxRight, startPos + rightSteps);
    return prefix[r + 1] - prefix[l];
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Solution:
  def maxTotalFruits(
      self,
      fruits: list[list[int]],
      startPos: int,
      k: int,
  ) -> int:
    ans = 0
    maxRight = max(startPos, fruits[-1][0])
    amounts = [0] * (1 + maxRight)
    for position, amount in fruits:
      amounts[position] = amount
    prefix = list(itertools.accumulate(amounts, initial=0))

    def getFruits(leftSteps: int, rightSteps: int) -> int:
      l = max(0, startPos - leftSteps)
      r = min(maxRight, startPos + rightSteps)
      return prefix[r + 1] - prefix[l]

    # Go right first.
    for rightSteps in range(min(maxRight - startPos, k) + 1):
      leftSteps = max(0, k - 2 * rightSteps)  # Turn left
      ans = max(ans, getFruits(leftSteps, rightSteps))

    # Go left first.
    for leftSteps in range(min(startPos, k) + 1):
      rightSteps = max(0, k - 2 * leftSteps)  # Turn right
      ans = max(ans, getFruits(leftSteps, rightSteps))

    return ans