โปรแกรมหาจำนวนเด็กที่ได้ลูกอม ตามกฎการแจกใน Python

กำหนด k ลูกอม แจกให้เด็กตามกฎ:

• เด็กคนที่ i ได้ i^2 ลูกอม
• เด็กคนที่ i ไม่ได้ลูกอมจนกว่าเด็ก 1 ถึง i-1 จะได้รับครบ
• ถ้าเด็กคนที่ i ไม่ได้ i^2 ลูกอม ถือว่าแจกไม่ถูกต้อง

เช่น k = 20 ผลลัพธ์คือ 3 เพราะ คนที่ 1 ได้ 1 คนที่ 2 ได้ 4 คนที่ 3 ได้ 9 คนที่ 4 ต้องการ 16 แต่เหลือแค่ 6 จึงไม่ถูกต้อง

วิธีแก้ปัญหา

• left := 0, right := k
• while right - left > 1
  • mid := floor((left + right) / 2)
  • ถ้า floor(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))

Input

20

Output

3