Problem Statement
leetcode problem link
Brute Force [TLE]
class Solution:
def distributeCandies(self, n: int, limit: int) -> int:
res = 0
cache = defaultdict(int)
def dfs(start, child, remain):
if (start, child, remain) in cache:
return cache[(start, child, remain)]
if child == 2:
if remain == 0:
return 1
return 0
ans = 0
for j in range(limit + 1):
ans += dfs(j, child + 1, remain - j)
cache[(start, child, remain)] = ans
return ans
for i in range(limit + 1):
res += dfs(i, 0, n - i)
return res
Editorial
Approach 1: Enumeration
class Solution:
def distributeCandies(self, n: int, limit: int) -> int:
ans = 0
for i in range(min(limit, n) + 1):
if n - i > 2 * limit:
continue
ans += min(n - i, limit) - max(0, n - i - limit) + 1
return ans
Approach 2: Inclusion-Exclusion Principle
def cal(x):
if x < 0:
return 0
return x * (x - 1) // 2
class Solution:
def distributeCandies(self, n: int, limit: int) -> int:
return (
cal(n + 2)
- 3 * cal(n - limit + 1)
+ 3 * cal(n - (limit + 1) * 2 + 2)
- cal(n - 3 * (limit + 1) + 2)
)
