Candy

n명의 아이들이 한 줄로 서있다. 각각의 아이에게는 평가 점수가 매겨져 있는데 이 정보는 ratings 배열에 들어있다.

다음 조건을 만족하도록 아이에게 사탕을 주려고 한다:

• 각각의 아이는 최소 하나의 사탕은 받아야 한다.
• 주변의 아이보다 더 높은 평가 점수를 가진 아이는 더 많은 사탕을 가져야 한다.

아이들에게 나눠줄 사탕의 최소 개수를 구하자.

평가 점수 배열의 크기는 \(1 \sim 2 \times 10^4\) 이고 각 평가 점수 범위는 \(0 \sim 2 \times 10^4\) 이다.

예를 들어, ratings = [1, 0, 2]라고 하자. [2, 1, 2]로 나눠주는 것 최선이기 때문에 정답은 5이다.

시뮬레이션 해보기

배열을 가지고 뭔가 정렬을 하거나 보조 자료 구조를 쓰거나 하는 방법을 궁리해봤는데, 이 문제는 그런 게 아니라 일단 조건에 맞게 사탕을 잘 나눠주는 것이 중요한 것 같다. 시뮬레이션을 잘 해보자.

먼저 주변의 아이를 확인하려면, i번째 아이에 대해서 i-1과 i+1 모두를 확인해야 한다. 하지만 이걸 제대로 하려면 인덱스에 대한 엣지 케이스를 세심히 처리해줘야 하고, 또 한 방향으로 진행할 때 이전 방향을 다시 봐야하기 때문에 꽤 까다롭다. 좀더 깔끔한 방법은 배열을 정방향으로 한번 역방향으로 한번, 총 두 번 보는 것이다.

정방향으로 나눠주는 걸 먼저 생각해보자. 최소 하나는 나눠준 상태라고 하면, i 번째 아이의 점수가 i-1 번째 아이보다 큰 경우에는 i-1 번째 아이에게 나눠준 것보다 하나 더 주기만 하면 된다. 즉, 이웃한 아이 중 왼쪽(i-1)의 아이만 먼저 기준으로 살펴본 것이다.

그럼 역방향의 경우는 어떻게 하면 될까? 정방향에서 왼쪽 아이를 살펴봤으니 여기서는 오른쪽 아이를 마저 살펴봐야 한다. 오른쪽(i+1)아이보다 지금 아이가 사탕이 더 많으면, 이전처럼 그냥 오른쪽 아이보다 사탕을 하나 더 주면 될까? 정방향을 살펴보면서 이미 왼쪽 아이를 기준으로 조건을 맞춰놨기 때문에, 오른쪽 아이보다 하나 더 주는 것이 오히려 덜 주는 경우가 생길 수 있다. 따라서 둘 중 더 많은 것을 취해야 양쪽 아이 모두보다 많은 조건을 만족할 수 있다. 예를 들어 점수가 [1, 3, 4, 5, 2]인 경우 정방향으로 나눠주고 나면 [1, 2, 3, 4, 1]이 되는데, 역방향으로 나눠줄 때 무작정 i+1 아이보다 하나 더 줘버리면 오히려 [1, 2, 3, 2, 1]로 4번째 아이의 사탕 개수가 줄어들어버릴 수 있기 때문이다.

이 아이디어를 구현하면 다음과 같다.

def candy(ratings):
    n = len(ratings)
    candies = [1] * n
    for i in range(1, n):
        if ratings[i] > ratings[i-1]:
            candies[i] = candies[i-1] + 1
    for i in range(n-2, -1, -1):
        if ratings[i] > ratings[i+1]:
            candies[i] = max(candies[i], candies[i+1] + 1)
    return sum(candies)

시간과 공간 복잡도 모두 O(N)이다.