aboutsummaryrefslogtreecommitdiff
path: root/distribute_candies_children2.py
blob: 6123df88e4d110ed6e2b4c7a6389492161b6ce05 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# 2929. Distribute Candies Among Children II
from math import comb

"""
you are given two positive integers 'n' and 'limit'. return the total number
of ways to distritube 'n' candies among 3 children such that no child gets
more than 'limit' candies.
"""


class Solution(object):
    def distributeCandies(self, n, limit):
        """
        :type n: int
        :type limit: int
        :rtype: int
        """
        ans = (n + 2) * (n + 1) // 2
        for i in range(1, 4):
            if n - i * (limit + 1) < 0:
                break
            ans += (
                (-1) ** i
                * comb(3, i)
                * (n - i * (limit + 1) + 2)
                * (n - i * (limit + 1) + 1)
                // 2
            )
            return ans


if __name__ == "__main__":
    obj = Solution()
    print(obj.distributeCandies(n=5, limit=2))
    print(obj.distributeCandies(n=3, limit=3))