Three different Python solutions


  • 0

    When I first saw this problem, a naive greedy method came into my mind. Let's set the initial condition that every child has 1 candy and update the candies from the child with smallest rating value. This one runs in O(nlog(n)) time and needs O(n) space.

    class Solution(object):
        def candy(self, ratings):
            """
            :type ratings: List[int]
            :rtype: int
            """
            n = len(ratings)
            child = [1] * n
            for r, i in sorted((r, i) for i, r in enumerate(ratings)):
                lower = [child[j]+1 for j in (i-1, i+1) if -1 < j < n and ratings[j] < r]
                child[i] = max(lower or [1])
            return sum(child)
    

    And when the distribution graph shows it only beats 12.5% among all submission, I know I must miss something. Then, I came up with another solution via two passes. This one runs in O(n) time and needs O(n) space.

    class Solution(object):
        def candy(self, ratings):
            """
            :type ratings: List[int]
            :rtype: int
            """
            n = len(ratings)
            child = [1] * n
            for i in range(1, n):
                if ratings[i] > ratings[i-1]: child[i] = max(child[i], child[i-1] + 1)
            for i in range(n-1)[::-1]:
                if ratings[i] > ratings[i+1]: child[i] = max(child[i], child[i+1] + 1)
            return sum(child)
    

    And finally, inspired by @allenlipeng47 's blog here, an O(n) time and O(1) space version and this one beats 98%.

    class Solution(object):
        def candy(self, ratings):
            """
            :type ratings: List[int]
            :rtype: int
            """
            n = len(ratings)
            ans = down = height = 0
            for i in range(n):
                if not i:
                    ans = prev = 1
                elif ratings[i] >= ratings[i-1]:
                    if down:
                        ans += (down+1) * down // 2 + max(down + 1 - height, 0)
                        down = height = 0
                    prev = prev + 1 if ratings[i] > ratings[i-1] else 1
                    ans += prev
                else:
                    if not height: height = prev
                    down += 1
                    prev = 1
            if down: ans += max(down + 1 - height, 0)
            return ans + (down+1) * down // 2
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.