Minimum number of candies needed for distribution

minimum_candies.py

class Solution(object):

    def solve(self, a):
        candies = [1 for i in a]
        n = len(a)

        # first pass
        for i in range(1, n):
            if a[i] > a[i-1]:
                candies[i] = candies[i-1] + 1

        # second pass
        for i in range(n-2, -1, -1):
            if a[i] > a[i+1]:
                candies[i] = max(candies[i+1] + 1, candies[i])

        return sum(candies)


if __name__ == "__main__":
    a = [5, 5, 4, 3, 2, 1]
    ans = Solution().solve(a)
    print(ans)