Leetcode: 135. Candy

Can we pre-process the input in a way to make easy to solve the problem? Sort a list of pairs \((v, i)\) in increasing order where \(ratings[i]=v\). For each child in order, assign its candy as describe in the rules. Time complexity is \(O(n \log n)\) and space is \(O(n)\).

from typing import List


class Solution:
    def candy(self, r: List[int]) -> int:
        N = len(r)
        low = [1] * N

        ans = 0
        for v, i in sorted(zip(r, range(N))):
            c = []
            if i - 1 >= 0 and r[i - 1] < v:
                c.append(low[i - 1] + 1)
            if i + 1 < N and r[i + 1] < v:
                c.append(low[i + 1] + 1)
            if c:
                low[i] = max(c)
            ans += low[i]

        return ans


assert Solution().candy([1, 0, 2]) == 5
assert Solution().candy([1, 2, 2]) == 4

Can we formulate the problem as a line chart? We want to assign \(1+2+..+k\) when we see a spike and \(p+(p-1)+...+1\) when we see a drop. If a spike is followed by a drop, we want to get \(max(k, p)\) for that child. Time and space complexity is \(O(n)\).

from typing import List


class Solution:
    def candy(self, r: List[int]) -> int:
        N = len(r)

        cand = [1] * N
        for i in range(1, N):
            if r[i] > r[i - 1]:
                cand[i] = cand[i - 1] + 1

        for i in range(N - 2, -1, -1):
            if r[i] > r[i + 1]:
                cand[i] = max(cand[i], cand[i + 1] + 1)

        return sum(cand)


assert Solution().candy([1, 0, 2]) == 5
assert Solution().candy([1, 2, 2]) == 4