Balls in urns with different max capacity for each urn

Question

I'm looking to figure out the formula for putting K balls in N urns, where each of the N urns has a specific max capacity. For example,

• Urn A has capacity for at most 5 balls
• Urn B has capacity for at most 2 balls
• Urn C has capacity for at most 1 ball
• Urn D has capacity for at most 2 balls

How many ways are there to put 5 indistinguishable balls in urns A, B, C and D?

With brute-force code, I know that this example has these 18 solutions:

(0, 2, 1, 2) (1, 1, 1, 2) (1, 2, 0, 2) (1, 2, 1, 1) (2, 0, 1, 2) (2, 1, 0, 2) (2, 1, 1, 1) (2, 2, 0, 1) (2, 2, 1, 0) (3, 0, 0, 2) (3, 0, 1, 1) (3, 1, 0, 1) (3, 1, 1, 0) (3, 2, 0, 0) (4, 0, 0, 1) (4, 0, 1, 0) (4, 1, 0, 0) (5, 0, 0, 0)

But I'm having trouble coming up with a general formula, which I need for my application, where the numbers get too large to brute-force like this. Does anyone know how to compute the number of combinations for putting indistinguishable balls in urns, only with max capacity for each urn?

Answer 1

After researching, and in particular looking at the answers to this post Number of ways to distribute indistinguishable balls into distinguishable boxes of given size and this post Using generating functions in combinatorics, I realize that @MathLover's comment above is right — thank you! Quoted here:

For generating function, if an urn can have max of 1 ball, write it as $(_0+_1)$ or $(1+)$, if it can have at most 2 balls write it as $(1++_2)$ and so on... So it comes down to finding coefficient of $_5$ in $(1+)(1++_2)^2(1++_2+_3+_4+_5)$

For those less interested in the pure formula than in the implementation, as I am, here's the code I'm now using:

import numpy as np

def comb_r_under_max(r,max_by_urn):
  mul = np.poly1d(np.ones(max_by_urn[0]+1)) # initializing polynomial
  for deg in max_by_urn[1:]:
    mul = np.polymul(np.poly1d(np.ones(deg+1)), mul) # multiplying
  return int(mul[r])

  
# test on our example
max_by_urn = [5,2,1,2]
r = 5

print(comb_r_under_max(r,max_by_urn)) # prints 18

Answer 2

one idea is to solve this kind of problems with recusion and lru_cache, from ground up, code can be simplified but I am sure you can work it out

@lru_cache
def get_combs(no_balls: int, urns: tuple) -> int:
    if no_balls == 0:
        return 1
    if no_balls > sum(urns):
        return 0
    if len(urns) ==  1:
        if urns[0]>=no_balls:
            return 1
        else:
            return 0
    res = 0
    for i in range(min(urns[0], no_balls)+1):
        res += get_combs(no_balls-i, urns[1:])
    return res