แจกลูกอมให้คนใน Python

แจกลูกอมให้ n คน แถวเดียว: คนแรก 1 ลูก คนที่ 2 ได้ 2 ลูก ... คนที่ n ได้ n ลูก แล้ววนกลับมาคนแรกได้ n+1 ลูก คนที่ 2 ได้ n+2 ลูก ... จนลูกอมหมด คนสุดท้ายได้ลูกอมที่เหลือทั้งหมด

ตัวอย่าง: candies=7, n=3 ผลลัพธ์ [2,2,3] (รอบ1: 1,2,3 รอบ2: คนแรกได้อีก 1)

ขั้นตอนการแก้ปัญหา

1. สร้างอาร์เรย์ res ขนาด n เติม 0
2. index = 0
3. วนลูปจน candies > 0:
   • res[index % n] += min(candies, index + 1)
   • candies -= (index + 1)
   • index += 1
4. คืน res

ตัวอย่างโค้ด

class Solution(object):
    def distributeCandies(self, candies, num_people):
        res = [0 for i in range(num_people)]
        index = 0
        while candies > 0:
            res[index % num_people] += min(candies, index + 1)
            candies -= (index + 1)
            index += 1
        return res

ob1 = Solution()
print(ob1.distributeCandies(8, 3))

อินพุต

8 3

เอาต์พุต

[3, 2, 3]

วิธีแก้แบบคณิตศาสตร์ (เร็วกว่า)

import math

def distributeCandies(candies, num_people):
    res = [0] * num_people
    # หาจำนวนรอบเต็ม k ที่แจกได้
    k = int((math.sqrt(1 + 8 * candies) - 1) // 2)
    # แจกรอบเต็ม
    for i in range(num_people):
        full_rounds = k // num_people
        extra = 1 if (k % num_people) > i else 0
        times = full_rounds + extra
        first = i + 1
        last = first + (times - 1) * num_people
        res[i] = times * (first + last) // 2
    # แจกลูกอมเหลือ
    remaining = candies - k * (k + 1) // 2
    res[k % num_people] += remaining
    return res