Problem
Solution 1: Range sum query
Assuming we can collect fruits in range [l, r], we need a fast query to compute the sum of those fruits.
Given startPos and k, we have four options:
1. move i steps to the left
2. move i steps to the left and k – i steps to the right.
3. move i steps to the right
4. move i steps to the right and k – i steps to the left.
We enumerate i steps and calculate maximum range [l, r] covered by each option, and collect all the fruit in that range.
Time complexity: O(m + k)
Space complexity: O(m)
where m = max(max(pos), startPos)
C++
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 |
// Author: Huahua class Solution { public: int maxTotalFruits(vector<vector<int>>& fruits, int startPos, int k) { const int m = max(startPos, (*max_element(begin(fruits), end(fruits)))[0]); vector<int> sums(m + 2); for (int i = 0, j = 0; i <= m; ++i) { sums[i + 1] += sums[i]; while (j < fruits.size() && fruits[j][0] == i) sums[i + 1] += fruits[j++][1]; } int ans = 0; for (int s = 0; s <= k; ++s) { if (startPos - s >= 0) { int l = startPos - s; int r = min(max(startPos, l + (k - s)), m); ans = max(ans, sums[r + 1] - sums[l]); } if (startPos + s <= m) { int r = startPos + s; int l = max(0, min(startPos, r - (k - s))); ans = max(ans, sums[r + 1] - sums[l]); } } return ans; } }; |
Solution 2: Sliding Window
Maintain a window [l, r] such that the steps to cover [l, r] from startPos is less or equal to k.
Time complexity: O(n)
Space complexity: O(1)
C++
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
// Author: Huahua class Solution { public: int maxTotalFruits(vector<vector<int>>& fruits, int startPos, int k) { auto steps = [&](int l, int r) { if (r <= startPos) return startPos - l; else if (l >= startPos) return r - startPos; else return min(startPos + r - 2 * l, 2 * r - startPos - l); }; int ans = 0; for (int r = 0, l = 0, cur = 0; r < fruits.size(); ++r) { cur += fruits[r][1]; while (l <= r && steps(fruits[l][0], fruits[r][0]) > k) cur -= fruits[l++][1]; ans = max(ans, cur); } return ans; } }; |