Number of ways to pick at least x balls (balls are identical)? - python-2.7

Question

There are N distinct boxes of balls in total. There are P boxes each containing A number of balls and remaining Q boxes contains B number of balls each.
Given a number X, what are the total the number of ways in which you can pick at least X balls from the boxes.
P+Q = N
Example: Number of P boxes=2 which contain 2 balls each, Number of Q boxes=1 which contain 2 balls. X=3(Given) where x=minimum number of balls to be picked
So, P+Q=3 (total number of boxes)
Combinations for the number of ways to pick atleast x i.e. 3 balls would be:
combinations of 3:(111),(210),(120),(021),(012),(201),(102)
combinations of 4:(220)(202)(022)(211)(121)(112)
combinations of 5:(212)(122)(221)
combinations of 6: (222)
total Combinations: 17
My Approach:
I have used "Stars and Bars Approach":
To calculate combinations of 6: x+y+z=6 which is converted into (2-x)+(2-y)+(2-z)=6 giving out x+y+z=0.
So, the combination of 6 becomes Binomial(2C2)=1
Similarly, Combinations of 5 becomes Binomial(3C2)=3
Combinations of 4= Binomial(4C2)=6
Combinations of 3= Binomial(5C2)=10
1+3+6+10=20
but the answer should be 1+3+6+7=17
Edge case has appeared on calculating the combinations of 3. How should I tackle this problem?
EDIT: CODE ADDED in python
global total_combinations
total_combinations=0
from math import factorial
def combinations(a):
global total_combinations
bars=numberofAs+numberofBs-1
stars=a
total_combinations+=factorial(stars+bars)/(factorial(bars)*factorial(stars))
numberofAs,numberofBs,numberofballsinA,numberofballsinB=map(int,raw_input().split())
x=int(raw_input())
operational_array=[]
for i in range(numberofAs):
operational_array.append(numberofballsinA)
for i in range(numberofBs):
operational_array.append(numberofballsinB)
max_x=sum(operational_array) #calculate combinations from x to max_x
k=max_x
for i in range(max_x,x-1,-1):
k=max_x-i
combinations(k)
print total_combinations

Answer 1

The number of balls you can take out of a box containing A balls is the number of balls you can put into an empty box of capacity A.
There are known formulas for that problem.
If all boxes have the same number of balls initially (as in the given example, in which A=B=2),
and that number is equal to or greater than the total number of balls to be removed from the boxes, then "stars and bars" will work.
But if the number of balls to be removed is greater than the number in a single box, there is an iterative formula to find the number of ways the balls can be selected.
To remove t balls from k boxes containing m balls each,
from scipy.special import comb
def combinations_with_limit(t, k, m):
total = 0
max_full_boxes = min(k, int(t/(m + 1)))
for i in range(max_full_boxes + 1):
total += int((-1)**i) * comb(k, i, exact=True) * comb(t + k - 1 - i*(m + 1), k - 1, exact=True)
return total
This is based on the formula in this math.stackexchange answer, but using t rather than n for the total number of balls removed in order to avoid confusion with the use of N in this question.
You can optimize and improve the style of this code, of course
(for example, I wouldn't suggest writing int((-1)**i) in production code);
the reason it's written this way is to stay as close as practical to the format of the MSE answer.
Not surprisingly, we have to think a little harder in the case where A and B are different.
To remove a total of t balls from p boxes containing a balls each and q boxes containing b balls each,
def combinations_with_two_limits(t, p, q, a, b):
total = 0
min_balls_from_p = max(0, t - q*b)
max_balls_from_p = min(t, p*a)
for i in range(min_balls_from_p, max_balls_from_p + 1):
total += combinations_with_limit(i, p, a) * combinations_with_limit(t - i, q, b)
return total
The idea here is that you first decide how to allocate the t balls into two groups, one to be removed from the boxes containing a balls and the other to be removed from the boxes containing b balls,
and then count all the ways you can select those subsets of the balls from those subsets of the boxes.
It may be possible to optimize the code further by going back to the derivation of the MSE formula cited above (either through generating functions or the inclusion-exclusion principle), but I wouldn't try it unless it's really critical to shave a few percentage points off the running time.
To remove at least X balls from the boxes, take the sum of the values of
combinations_with_two_limits as t takes on all integer values from X up to and including the largest number of balls you can remove
(which is P*A + Q*B).

Answer 2

Break it in 2 parts where you pick k balls from the P boxes of capacity A and x-k balls from the Q bins with capacity B. Use stars and bars approach on each side to calculate the number of ways to pick k,x-k respectively, and multiply to get the total. (Considering of course that you can only do that when both k,x-k do not exceed the capacities A,B).
To not consider the cases where k>A and x-k>B you have to implement the formula below
Which counts all combinations in which n items are allocated to k bins, excluding combinations where at least one bin contains more than C items.
In your case you need to use this formula twice, once to allocate k balls to P bins with capacity A, and once to allocate x-k balls to Q bins with capacity B.
Then the multiplication should give you the correct result.