Chương trình đếm số trẻ nhận kẹo theo quy tắc trong Python

Cho k kẹo. Trẻ thứ i nhận i^2 kẹo. Phải phân phát theo thứ tự, thiếu kẹo thì dừng. Tìm số trẻ được phát kẹo.

Thuật toán

• Tổng kẹo cho m trẻ đầu: S(m) = m(m+1)(2m+1)/6.
• Tìm m lớn nhất thỏa S(m) ≤ k bằng tìm kiếm nhị phân.

Ví dụ

def solve(k):
    left, right = 0, 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

print(solve(20))  # Output: 3

Giải thích

k=20: 1^2+2^2+3^2=14 ≤20, 4^2=16 nhưng chỉ còn 6 kẹo → 3 trẻ.