Distribute Candies Among Children II

Tested Python solution for LeetCode 2929 with 18 pytest cases. Generate a practice environment with lcpy.

LeetCode 2929, Medium. Topics: Math, Combinatorics, Enumeration. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 18 parametrized pytest cases, and a playground notebook:

lcpy gen -n 2929   # by problem number
lcpy gen -s distribute_candies_among_children_ii   # by problem name

​

Problem

You are given two positive integers n and limit.

Return the total number of ways to distribute n candies among 3 children such that no child gets more than limit candies.

​

Examples

Input: n = 5, limit = 2
Output: 3
Explanation: There are 3 ways to distribute 5 candies such that no child gets more than 2 candies: (1, 2, 2), (2, 1, 2) and (2, 2, 1).

Input: n = 3, limit = 3
Output: 10
Explanation: There are 10 ways to distribute 3 candies such that no child gets more than 3 candies: (0, 0, 3), (0, 1, 2), (0, 2, 1), (0, 3, 0), (1, 0, 2), (1, 1, 1), (1, 2, 0), (2, 0, 1), (2, 1, 0) and (3, 0, 0).

​

Constraints

• 1 <= n <= 10^6
• 1 <= limit <= 10^6

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

class Solution:
    # Time: O(1)
    # Space: O(1)
    def distribute_candies(self, n: int, limit: int) -> int:
        def capped_ways(total: int) -> int:
            # Ways for 3 non-negative parts summing to `total`, ignoring the cap:
            # C(total + 2, 2); 0 when no composition exists.
            return (total + 2) * (total + 1) // 2 if total >= 0 else 0

        return (
            capped_ways(n)
            - 3 * capped_ways(n - (limit + 1))
            + 3 * capped_ways(n - 2 * (limit + 1))
            - capped_ways(n - 3 * (limit + 1))
        )

​

Complexity

Time	Space
O(1)	O(1)

​

Tags

NeetCode All.