โปรแกรมตรวจสอบว่ากินขนมชนิดโปรดในวันที่โปรดได้หรือไม่ใน Python

ปัญหา

มี candiesCount[i] = จำนวนขนมชนิด i และ queries[i] = [favoriteType, favoriteDay, dailyCap] กฎ: เริ่มกินวันที่ 0, ต้องกินขนมชนิดก่อนหน้าครบก่อนถึงชนิดถัดไป, ต้องกินอย่างน้อย 1 ชิ้นต่อวัน จนกว่าจะหมด ตอบแต่ละ query ว่าสามารถกินขนมชนิด favoriteType ในวัน favoriteDay ได้หรือไม่ โดยไม่กินเกิน dailyCap ต่อวัน

ตัวอย่าง

candiesCount = [7,4,5,3,8], queries = [[0,2,2],[4,2,4],[2,13,100]] → [true, false, true]

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

1. sumcandy := [candiesCount[0]] (prefix sum)
2. สร้าง prefix sum ทั้งหมด แล้วเพิ่ม 0 ท้ายสุด
3. สำหรับแต่ละ query:
   • typ, day, cap := query
   • ถ้า day+1 > sumcandy[typ] หรือ (day+1)*cap <= sumcandy[typ-1]: False
   • มิฉะนั้น: True

การนำไปใช้งาน (Python)

def solve(candiesCount, queries):
    sumcandy = [candiesCount[0]]
    index = 1
    while index < len(candiesCount):
        sumcandy.append(sumcandy[index-1] + candiesCount[index])
        index += 1
    sumcandy.append(0)
    res = []
    for each in queries:
        typ = each[0]
        day = each[1]
        cap = each[2]
        if day+1 > sumcandy[typ] or (day+1)*cap <= sumcandy[typ-1]:
            res.append(False)
        else:
            res.append(True)
    return res

candiesCount = [7,4,5,3,8]
queries = [[0,2,2],[4,2,4],[2,13,100]]
print(solve(candiesCount, queries))

อินพุต

[7,4,5,3,8], [[0,2,2],[4,2,4],[2,13,100]]

เอาต์พุต

[True, False, True]