aboutsummaryrefslogtreecommitdiff
path: root/max_fruits_harvested_k_step.py
diff options
context:
space:
mode:
Diffstat (limited to 'max_fruits_harvested_k_step.py')
-rw-r--r--max_fruits_harvested_k_step.py72
1 files changed, 72 insertions, 0 deletions
diff --git a/max_fruits_harvested_k_step.py b/max_fruits_harvested_k_step.py
new file mode 100644
index 0000000..bc5b1b0
--- /dev/null
+++ b/max_fruits_harvested_k_step.py
@@ -0,0 +1,72 @@
+# 2106. Maximum Fruits Harvested After at Most K Steps
+import bisect
+
+"""
+fruits are available at some positions on an infinite x-axis. you are given a
+2d integer array 'fruits' where 'fruits[i] = [positions[i], amount[i]]'
+depicts 'amount[i]' fruits at the position 'position[i]'. 'fruits' is already
+sorted by 'position[i]' in ascending order, and eash 'position[i]' is unique.
+you are also given an integer 'startPos' and an integer 'k'. initially, y-u
+are at the position 'startPos'. from any position, you can either walk to the
+left or right. it takes one step to move one unit on the x-axis, and you can
+walk at most 'k' steps in total. for every position you reach, you harvest
+all the fruits at that position, and the fruits will dsiappear from that
+position.
+"""
+
+
+class Solution(object):
+ def maxTotalFruits(self, fruits, startPos, k):
+ """
+ :type fruits: List[List[int]]
+ :type startPos: int
+ :type k: int
+ :rtype: int
+ """
+ amt = {}
+ for a, b in fruits:
+ amt[a] = b
+ pos = [a for a, b in fruits if a != startPos]
+ left, right, n = [], [], len(pos)
+ idx = bisect.insort_right(pos, startPos)
+ curr_f = 0
+ for i in range(idx, n):
+ curr_pos = pos[i]
+ curr_f += amt[curr_pos]
+ right.append([curr_pos - pos, curr_f])
+ curr_f = 0
+ for i in range(idx - 1, -1, -1):
+ curr_pos = pos[i]
+ curr_f += amt[curr_pos]
+ left.append([pos - curr_pos, curr_f])
+ ans = 0
+ for r_dist, r_f in right:
+ if r_dist <= k:
+ curr_f = r_f
+ l_dist = k - 2 * r_dist
+ if l_dist > 0:
+ idx = bisect.bisect_right(left, [l_dist, float("inf")])
+ if idx > 0:
+ curr_f += left[idx - 1][1]
+ ans = max(ans, curr_f)
+ for l_dist, l_f in left:
+ if l_dist <= k:
+ curr_f = l_f
+ r_dist = k - 2 * l_dist
+ if r_dist > 0:
+ idx = bisect.bisect_right(right, [r_dist, float("inf")])
+ if idx > 0:
+ curr_f += right[idx - 1][1]
+ ans = max(ans, curr_f)
+ return ans + amt.get(pos, 0)
+
+
+if __name__ == "__main__":
+ obj = Solution()
+ print(obj.maxTotalFruits(fruits=[[2, 8], [6, 3], [8, 6]], startPos=5, k=4))
+ print(
+ obj.maxTotalFruits(
+ fruits=[[0, 9], [4, 1], [5, 7], [6, 2], [7, 4], [10, 9]], startPos=5, k=4
+ )
+ )
+ print(obj.maxTotalFruits(fruits=[[0, 3], [6, 4], [8, 5]], startPos=3, k=2))