규칙을 지키며 사탕을 아이들에게 분배하는 파이썬 프로그램

문제 개요

k개의 사탕이 주어졌을 때, 이를 여러 아이들에게 나누어 주는 문제입니다. 단, 반드시 지켜야 할 규칙이 있습니다.

i번째 아이는 정확히 i²개의 사탕을 받아야 합니다.
i번째 아이에게 사탕을 주기 전에, 1번부터 i-1번까지의 모든 아이가 먼저 사탕을 받아야 합니다(순서대로 지급).
i번째 아이가 i²개를 온전히 받지 못하면, 그 지급은 유효하지 않습니다.

예시

예를 들어 k = 20이라고 가정해 보겠습니다. 첫 번째 아이는 1개, 두 번째 아이는 2² = 4개, 세 번째 아이는 3² = 9개를 받습니다. 네 번째 아이는 4² = 16개가 필요하지만 남은 사탕은 6개뿐이므로 유효한 지급이 불가능합니다. 따라서 정답은 3이 됩니다.

접근 방법: 이진 탐색

1부터 n까지 제곱수의 합은 잘 알려진 공식 n(n+1)(2n+1)/6으로 구할 수 있습니다. 이 공식을 활용하면, k개의 사탕으로 최대 몇 명까지 지급할 수 있는지 이진 탐색으로 빠르게 찾을 수 있습니다.

풀이 절차는 다음과 같습니다.

left := 0, right := k 로 초기화합니다.
right - left > 1 인 동안 다음을 반복합니다.
- mid := (left + right) / 2 의 내림값
- mid × (mid + 1) × (2 × mid + 1) / 6 의 내림값이 k보다 크면 right := mid, 그렇지 않으면 left := mid
반복 종료 후, right × (right + 1) × (2 × right + 1) ≤ k × 6 이면 right를 반환하고, 그렇지 않으면 left를 반환합니다.

파이썬 구현 예제

def solve(k):
   left = 0
   right = k
   while (right - left > 1):
      mid = (left + right) // 2
      if (mid * (mid + 1) * (2 * mid + 1) // 6 > k):
         right = mid
      else:
         left = mid
   if (right * (right + 1) * (2 * right + 1) <= k * 6):
      return right
   return left

k = 20
print(solve(k))

입력

출력

이 알고리즘의 시간 복잡도는 O(log k)입니다. 단순히 한 명씩 더해 가며 확인하는 O(√k) 방식보다도 효율적이며, k가 매우 큰 값이더라도 빠르게 정답을 구할 수 있다는 장점이 있습니다.