Combinatorial problem and the search for an index

Question

although I encountered this problem while coding, I think it is more math than coding.

The task is simple: I have three baskets of size $B_1$, $B_2$, and $B_3$ (every number is an integer here) and $A <= B_1 + B_2 + B_3$ units to distribute among those baskets. Basket 1 can hold between 0 and $B_1$ units, and so on. This gives me a list of possible combinations to distribute the A units, and this list is sorted (I can choose the sorting algorithm as I want to, and am currently working with the lexicographical order, i.e. the first is $(0,0,A)$, the second $(0,1,A-1)$, and so on).

The question: Having an element $(A_1,A_2,A_3)$ with $A_1 + A_2 + A_3 = A$ and knowing $B_1$, $B_2$, and $B_3$, is there a formula that tells me the index of $(A_1,A_2,A_3)$ in the list of possible combinations?

Example: If $(B_1,B_2,B_3) = (10,30,60)$ and $A = 61$, there are 340 combinations, ${(0,1,61),(0,2,60),...,(10,30,51)}$. The index of $(5,20,36)$ is then $174$.

Any ideas? As it is part of a larger code, I could always back up the formula with a mundane index search via a loop, so a formula that "only" works in 99% of the times would already be helpful. I found something that works in my example for every $A$ but those between 60 and 70, which is not enough to speed up my code.

Thanks a lot in advance :-)

PS: In my coding exercise, I can reduce the number of combinations significantly, but leaves the set of combinations without any foreseeable structure. Unfortunately, the index search still takes very long. My alternative would be to just not reduce them, work with the large set of combinations but utilize the structure of this set. However, I would need a short/fast formula for this then.

Answer 1

I will assume zero-indexing, since it happens to be more convenient. If you want the first triple to be indexed one, you will need to add one to this answer. This is a constant time solution.

The (zero-)index of a valid triple $(A_1,A_2,A_3)$ is... $$ (\text{# of valid triples $(C_1,C_2,C_3)$ with $C_1<A_1$})+ \\ (\text{# of valid triples $(C_1,C_2,C_3)$ with $C_1=A_1,\;C_2<A_2$})+ \\ (\text{# of valid triples $(C_1,C_2,C_3)$ with $C_1=A_1,\;C_2=A_2$},C_3<A_3) $$ This is because the index is equal to the number of triples which are below it lexicographically. This kind of analysis works in general when finding the index of a combinatorial object that can be represented as a sequence of integers, like permutations, subsets, multisets, etc.

Now we just need to solve these three counting problems. Starting with the first summand, the number of valid triples $(C_1,C_2,C_3)$ with $C_1<A_1$ is the number of nonnegative integer solutions to this system of equations and inequalities: $$ x_1+x_2+x_3=A,\\ \tag1 0\le x_1<A_1,\qquad 0\le x_2<B_2+1,\qquad 0\le x_3<B_3+1 $$ Using the methods in extended stars-and-bars problem(where the upper limit of the variable is bounded), you can show that the number of solutions to $(1)$ is $$ \binom{A+2}{2} -\binom{A+2-A_1}{2}^*-\binom{A+2-(B_2+1)}{2}^*-\binom{A+2-(B_3+1)}{2}^*\\ +\binom{A+2-(A_1+(B_2+1))}{2}^*+\binom{A+2-(A_1+(B_3+1))}{2}^*+\binom{A+2-(B_2+1+B_3+1)}{2}^* $$ The ${}^*$ refers to the following caveat: $$\binom{n}k^* \text{ is defined to be zero if $n<0$}$$

Next, you need to count valid triples of the form $(C_1,C_2,C_3)$ with $C_1=A_1$ and $C_2<A_2$. This time, we are counting integer solutions to $$ x_2+x_3=A-A_1\\ 0\le x_2 < A_2\qquad 0\le x_3 < B_3+1 \tag2 $$ which using the same method, is $$ \binom{A-A_1+1}{1}-\binom{A-A_1+1-A_2}{1}^*-\binom{A-A_1+1-(B_3+1)}{1}^* $$ With the ${}^*$ condition, this is equivalent to $$(A-A_1+1)-\max(A-A_1+1-A_2,0)-\max(A-A_1-B_3,0).$$ Finally, there are zero valid triples of the form $(A_1,A_2,C_3)$ with $C_3<A_3$, since such triples would sum to $A$. Therefore, you answer is just equal to the the sum of the two displayed formulae, which can be computed in pretty much $O(1)$ time.