Maximum Total Beauty of the Gardens LeetCode solution in Python

1
class Solution:
2
    def maximumBeauty(
3
        self, flowers: List[int], newFlowers: int, target: int, full: int, partial: int
4
    ) -> int:
5

6
        # move out already completed garden
7
        already_complete = 0
8
        temp = []
9
        for f in flowers:
10
            if f >= target:
11
                already_complete += 1
12
            else:
13
                temp.append(f)
14

15
        max_beauty = 0
16

17
        flowers = temp
18
        flowers.sort()
19

20
        presum = [0] + list(accumulate(flowers))
21
        N = len(flowers)
22

23
        dp_arr = []
24

25
        for i in range(N + 1):
26
            # iterate all index: left part is all partial, right part (>= i) is all complete
27

28
            # update the dp arr for binary search below
29
            if i < N:
30
                dp_arr.append((i + 1) * flowers[i] - presum[i + 1])
31

32
            right_sum = presum[-1] - presum[i]
33
            right_count = N - i
34

35
            # if can't make the right part all complete, go to next index
36
            if right_sum + newFlowers < right_count * target:
37
                continue
38

39
            # remaining flowers after making right part all complete
40
            rem_flowers = newFlowers - (right_count * target - right_sum)
41

42
            # binary search to find the maximum possible min flowers in the left part (named 'min_partial')
43
            if i == 0:
44
                min_partial = 0
45
            else:
46
                j = min(bisect.bisect_right(dp_arr, rem_flowers) - 1, i - 1)
47
                min_partial = min((rem_flowers + presum[j + 1]) // (j + 1), target - 1)
48

49
            complete = right_count + already_complete
50
            max_beauty = max(max_beauty, complete * full + min_partial * partial)
51

52
        return max_beauty