[Python] 31063: Candy Cane Feast

31063: Candy Cane Feast

USACO 2023 December Contest, Bronze
Problem 1. Candy Cane Feast

풀이

N * M처럼 보이는 간단한 구현 문제인데, O(N * M)으로는 풀 수 없어서 시간복잡도를 더 줄일 수 있는 방법이 있는지 한참 고민했다.

최악의 경우는 모든 소 N마리가 사탕 M개를 계속 나눠 먹는 상황이다. 이 경우 맨 앞의 소는 매번 사탕을 먹을 것이고, 현재 자신의 크기만큼 먹을 테니 키가 2배씩 커진다.
사탕 크기의 제한이 [1, 10^9]이므로, 맨 앞 소의 초기 키가 1이라고 가정하면: log2(10^9) = 약 30번 만에 맨 앞 소의 키가 사탕보다 커지게 된다. 30번 이후로는 맨 앞의 소가 사탕을 독식하게 된다.

맨 앞 소만 고려하더라도 연산 횟수가 많이 줄어든다. 사탕을 다 먹었을 경우 break하여야 시간 제한에 걸리지 않는다.

코드

import sys
input = lambda: sys.stdin.readline().strip()

n, m = list(map(int, input().split()))
cows = list(map(int, input().split()))
canes = list(map(int, input().split()))

for cane_top in canes:
    cane_bottom = 0
    for i in range(len(cows)):
        if cows[i] > cane_bottom:
            eat = min(cows[i] - cane_bottom, cane_top - cane_bottom)
            cows[i] += eat
            cane_bottom += eat
            if cane_top == cane_bottom:
                break 

print('\n'.join(map(str, cows)))