파이썬으로 특정 날짜에 원하는 캔디를 먹을 수 있는지 확인하는 프로그램

캔디 종류별 개수가 저장된 배열 candiesCount가 있다고 가정해 봅시다. 여기서 candiesCount[i]는 i번째 종류의 캔디가 몇 개 있는지를 나타냅니다. 또 다른 배열 queries도 주어지며, queries[i]는 [favoriteType_i, favoriteDay_i, dailyCap_i]라는 세 가지 값으로 구성됩니다.

문제의 규칙

0일차부터 캔디를 먹기 시작합니다.
i번째 종류의 캔디는 앞선 0번부터 i-1번까지의 모든 캔디를 다 먹은 후에야 먹을 수 있습니다.
모든 캔디를 소진할 때까지 매일 최소 한 개 이상의 캔디를 먹어야 합니다.

위 규칙을 지키면서 각 쿼리에 대한 결과를 불리언(Boolean) 값 배열로 만들어야 합니다. i번째 결과는 어떤 날에도 dailyCap_i개를 초과하지 않고 캔디를 먹는다는 조건 하에, favoriteDay_i일째 되는 날에 favoriteType_i 종류의 캔디를 먹을 수 있다면 true가 됩니다. 참고로 규칙 2에 따라 같은 날 서로 다른 종류의 캔디를 함께 먹는 것도 가능합니다.

예시 살펴보기

입력이 candiesCount = [7,4,5,3,8], queries = [[0,2,2],[4,2,4],[2,13,100]]이라면 출력은 [true, false, true]가 됩니다. 그 이유는 다음과 같습니다.

0일차와 1일차에 각각 2개씩 0번 종류 캔디를 먹으면, 2일차에도 여전히 0번 종류 캔디를 먹을 수 있습니다.
하루에 최대 4개씩 먹는 경우를 생각해 봅시다. 0일차에 0번 종류 캔디 4개를 먹고, 1일차에 남은 0번 종류 캔디와 1번 종류 캔디를 먹게 됩니다. 그러면 2일차에는 1번, 2번 종류 캔디까지만 도달할 수 있으므로, 2일차에 4번 종류 캔디를 먹는 것은 불가능합니다.
반면 하루에 1개씩만 먹으면, 13일차에 정확히 2번 종류 캔디를 먹게 됩니다.

해결 접근 방법

이 문제의 핵심은 누적합(prefix sum)입니다. 종류별 캔디 개수의 누적합 배열을 미리 계산해 두면, 각 쿼리를 O(1) 시간에 판정할 수 있습니다. 구체적인 단계는 다음과 같습니다.

첫 번째 요소가 candiesCount[0]인 누적합 리스트 sumcandy를 만듭니다.
인덱스 변수 index를 1로 초기화하고, candiesCount의 길이보다 작은 동안 이전 누적값과 현재 캔디 개수를 더해 sumcandy에 추가합니다.
sumcandy 끝에 0을 추가해 경계 조건을 처리합니다.
queries의 각 항목에 대해 typ(캔디 종류), day(날짜), cap(일일 상한)을 추출합니다.
다음 두 조건 중 하나라도 참이면 False를, 둘 다 아니면 True를 결과 리스트에 추가합니다.
- day+1 > sumcandy[typ] : 하루에 하나씩만 먹어도 해당 종류의 캔디가 day일 이전에 바닥나는 경우
- (day+1)*cap <= sumcandy[typ-1] : 하루에 최대치로 먹어도 day일까지 앞선 종류들의 캔디를 다 소진하지 못하는 경우
모든 쿼리를 처리한 뒤 결과 리스트를 반환합니다.

파이썬 구현 코드

아래 구현 예제를 통해 더 잘 이해해 봅시다.

def solve(candiesCount, queries):
    # 누적합 배열 생성
    sumcandy = [candiesCount[0]]
    index = 1
    while index < len(candiesCount):
        sumcandy.append(sumcandy[index-1] + candiesCount[index])
        index += 1
    sumcandy.append(0)

    res = []
    for each in queries:
        typ = each[0]
        day = each[1]
        cap = each[2]
        if day+1 > sumcandy[typ] or (day+1)*cap <= sumcandy[typ-1]:
            res.append(False)
        else:
            res.append(True)
    return res

candiesCount = [7,4,5,3,8]
queries = [[0,2,2],[4,2,4],[2,13,100]]
print(solve(candiesCount, queries))

입력

[7,4,5,3,8], [[0,2,2],[4,2,4],[2,13,100]]

출력

[True, False, True]

마무리 정리

이 풀이는 누적합 배열을 활용해 각 쿼리를 상수 시간에 처리합니다. "최소 섭취량(하루 1개)으로도 목표 종류에 도달할 수 있는가"와 "최대 섭취량(하루 cap개)으로도 목표 날짜까지 앞선 종류를 소진할 수 없는가"라는 두 가지 경계 조건만 확인하면 되기 때문에, 전체 시간 복잡도는 O(n + q)로 매우 효율적입니다. 여기서 n은 캔디 종류의 수, q는 쿼리의 개수입니다.