diff options
Diffstat (limited to 'max_fruits_harvested_k_step.py')
-rw-r--r-- | max_fruits_harvested_k_step.py | 72 |
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)) |