GeeksforGeeks App
Open App
Browser
Continue

# Count ways choosing K balls from any given A boxes

Given integers A and K, there are A boxes the first box contains K balls, the Second box contains K + 1 balls, the Third Box contains K + 3 balls so on till the A’th box contains K + A balls, the task for this problem is to count the number of ways of picking K balls from any of the boxes. (Print answer modulo 109 + 7).

Examples:

Input: A = 2, K = 2
Output: 4
Explanation: We have two boxes one with 2 balls let’s say Box 1 with Balls {A1, A2} and one with 3 balls let’s say Box 2 with balls {B1, B2, B3}

• First way: we can remove balls A1 and A2 from Box 1.
• Second way: we can remove balls B1 and B2 from Box 2.
• Third way: we can remove balls B2 and B3 from Box 2.
• Forth way : we can remove balls B1 and B3 from Box 2.

Input: A = 3, K = 3
Output: 15

Naive approach: The basic way to solve the problem is as follows:

The basic way to solve this problem is to generate all possible combinations by using a recursive approach.

Time Complexity: O(N!)
Auxiliary Space: O(1)

Efficient Approach:  The above approach can be optimized based on the following idea:

The problem can be solved with combinatorics:

• Total ways = KCK + K+1CK + K+2CK + ……. +  K+A-1CK
• we know N+1CM+1 = NCM + NCM+1 that is NCM = N+1CM+1NCM+1
• Total ways = K+1CK+1 + (K+2CK+1K+1CK+1) + (K+3CK+1K+2CK+1) + ……………. + (K+ACK+1K+A-1CK-1)
• Total ways = K+ACK+1
• Total ways = (K + A)! / ((K + 1)! * (A – 1)!)

Instead of dividing factorials of K + 1 and A – 1, we multiply their modular multiplicative inverses

Follow the steps below to solve the problem:

• Initializing fact[] array and Precomputing all factorials from 1 to 100000.
• initializing ANS variable.
• Calculating the answer by using the above formula.

Below is the implementation of the above approach:

## C++

 `// C++ code to implement the approach``#include ``using` `namespace` `std;` `// To avoid integer overflow``#define int long long` `const` `int` `MOD = 1e9 + 7;` `// To find GCD of a and b``int` `gcd(``int` `a, ``int` `b);` `// To compute x raised to power y``// under modulo M``int` `power(``int` `x, unsigned ``int` `y, unsigned ``int` `M);` `// Function to find modular inverse``// of a under modulo M``// Assumption: M is prime``int` `modInverse(``int` `A, ``int` `M)``{``    ``int` `g = gcd(A, M);``    ``return` `power(A, M - 2, M);``}` `// To compute x^y under modulo m``int` `power(``int` `x, unsigned ``int` `y, unsigned ``int` `M)``{``    ``if` `(y == 0)``        ``return` `1;` `    ``int` `p = power(x, y / 2, M) % M;``    ``p = (p * p) % M;` `    ``return` `(y % 2 == 0) ? p : (x * p) % M;``}` `// Function to return gcd of a and b``int` `gcd(``int` `a, ``int` `b)``{``    ``if` `(a == 0)``        ``return` `b;``    ``return` `gcd(b % a, a);``}` `// Function to Count number of ways``// balls can be removed from following``// boxes``int` `countWays(``int` `A, ``int` `K)``{` `    ``// Precomputation array storing``    ``// factorials MOD 1e9 + 7``    ``int` `fact[100001];` `    ``// Initialize the first element``    ``fact[1] = 1;` `    ``// Filling factorial table``    ``for` `(``int` `i = 2; i <= 100000; i++) {``        ``fact[i] = (fact[i - 1] * i) % MOD;``    ``}` `    ``// Storing Total number of arrangements``    ``int` `ans = (fact[A + K] * modInverse(fact[A - 1], MOD))``              ``% MOD;` `    ``// Multiplying inverse of (K + 1)!``    ``ans = (ans * modInverse(fact[K + 1], MOD)) % MOD;` `    ``// Returning the answer``    ``return` `ans;``}` `// Driver Code``int32_t main()``{` `    ``// Input 1``    ``int` `A = 2, K = 2;` `    ``// Function Call``    ``cout << countWays(A, K) << endl;` `    ``// Input 2``    ``int` `A1 = 3, K1 = 3;` `    ``// Function Call``    ``cout << countWays(A1, K1) << endl;``    ``return` `0;``}`

## Java

 `import` `java.math.BigInteger;` `public` `class` `Main``{``    ``// To avoid integer overflow``    ``private` `static` `final` `BigInteger MOD = BigInteger.valueOf(``1000000000` `+ ``7``);` `    ``// To find GCD of a and b``    ``private` `static` `BigInteger gcd(BigInteger a, BigInteger b) {``        ``if` `(a.equals(BigInteger.ZERO)) {``            ``return` `b;``        ``}``        ``return` `gcd(b.mod(a), a);``    ``}` `    ``// To compute x raised to power y under modulo M``    ``private` `static` `BigInteger power(BigInteger x, BigInteger y, BigInteger M) {``        ``if` `(y.equals(BigInteger.ZERO)) {``            ``return` `BigInteger.ONE;``        ``}` `        ``BigInteger p = power(x, y.divide(BigInteger.TWO), M).mod(M);``        ``p = p.multiply(p).mod(M);` `        ``return` `y.mod(BigInteger.TWO).equals(BigInteger.ZERO) ? p : x.multiply(p).mod(M);``    ``}` `    ``// Function to find modular inverse of a under modulo M``    ``private` `static` `BigInteger modInverse(BigInteger A, BigInteger M) {``        ``BigInteger g = gcd(A, M);``        ``return` `power(A, M.subtract(BigInteger.TWO), M);``    ``}` `    ``// Function to Count number of ways balls can be removed from following boxes``    ``private` `static` `BigInteger countWays(BigInteger A, BigInteger K)``    ``{``      ` `        ``// Precomputation array storing factorials MOD 1e9 + 7``        ``BigInteger[] fact = ``new` `BigInteger[``100001``];``        ``fact[``1``] = BigInteger.ONE;` `        ``// Filling factorial table``        ``for` `(``int` `i = ``2``; i <= ``100000``; i++) {``            ``fact[i] = fact[i - ``1``].multiply(BigInteger.valueOf(i)).mod(MOD);``        ``}` `        ``// Storing Total number of arrangements``        ``BigInteger ans = fact[A.intValue() + K.intValue()].multiply(modInverse(fact[A.intValue() - ``1``], MOD)).mod(MOD);` `        ``// Multiplying inverse of (K + 1)!``        ``ans = ans.multiply(modInverse(fact[K.intValue() + ``1``], MOD)).mod(MOD);` `        ``// Returning the answer``        ``return` `ans;``    ``}``  ` `    ``// Driver Code``    ``public` `static` `void` `main(String[] args) {``        ``// Input 1``        ``BigInteger A = BigInteger.valueOf(``2``);``        ``BigInteger K = BigInteger.valueOf(``2``);``        ``System.out.println(countWays(A, K));` `        ``// Input 2``        ``BigInteger A1 = BigInteger.valueOf(``3``);``        ``BigInteger K1 = BigInteger.valueOf(``3``);``        ``System.out.println(countWays(A1, K1));``    ``}``}` `// This code is contributed by ritaagarwal.`

## Python3

 `import` `math` `# To avoid integer overflow``MOD ``=` `int``(``1e9` `+` `7``)` `# To find GCD of a and b``def` `gcd(a: ``int``, b: ``int``) ``-``> ``int``:``    ``if` `a ``=``=` `0``:``        ``return` `b``    ``return` `gcd(b ``%` `a, a)` `# To compute x raised to power y under modulo M``def` `power(x: ``int``, y: ``int``, M: ``int``) ``-``> ``int``:``    ``if` `y ``=``=` `0``:``        ``return` `1` `    ``p ``=` `power(x, y ``/``/` `2``, M) ``%` `M``    ``p ``=` `(p ``*` `p) ``%` `M` `    ``return` `(p ``if` `y ``%` `2` `=``=` `0` `else` `x ``*` `p) ``%` `M` `# Function to find modular inverse of a under modulo M``# Assumption: M is prime``def` `modInverse(A: ``int``, M: ``int``) ``-``> ``int``:``    ``g ``=` `gcd(A, M)``    ``return` `power(A, M ``-` `2``, M)` `# Function to Count number of ways balls can be removed from following boxes``def` `countWays(A: ``int``, K: ``int``) ``-``> ``int``:``  ` `    ``# Precomputation array storing factorials MOD 1e9 + 7``    ``fact ``=` `[``0``] ``*` `(``100001``)``    ``fact[``1``] ``=` `1``    ``for` `i ``in` `range``(``2``, ``100001``):``        ``fact[i] ``=` `(fact[i ``-` `1``] ``*` `i) ``%` `MOD``        ` `    ``# Storing Total number of arrangements``    ``ans ``=` `(fact[A ``+` `K] ``*` `modInverse(fact[A ``-` `1``], MOD)) ``%` `MOD``    ` `    ``# Multiplying inverse of (K + 1)!``    ``ans ``=` `(ans ``*` `modInverse(fact[K ``+` `1``], MOD)) ``%` `MOD``    ` `    ``# Returning the answer``    ``return` `ans` `if` `__name__ ``=``=` `"__main__"``:``    ``# Input 1``    ``A ``=` `2``    ``K ``=` `2``    ` `    ``# Function Call``    ``print``(countWays(A, K))``    ` `    ``# Input 2``    ``A1 ``=` `3``    ``K1 ``=` `3``    ` `    ``# Function Call``    ``print``(countWays(A1, K1))` `# This code is contributed by lokeshpotta20.`

## C#

 `using` `System;``using` `System.Collections.Generic;` `class` `GFG``{``  ``static` `long` `MOD = 1000000007;` `  ``// Function to return gcd of a and b``  ``static` `long` `gcd(``long` `a, ``long` `b)``  ``{``    ``if` `(a == 0)``      ``return` `b;``    ``return` `gcd(b % a, a);``  ``}` `  ``// To compute x^y under modulo m``  ``static` `long` `power(``long` `x, ``long` `y, ``long` `M)``  ``{``    ``if` `(y == 0)``      ``return` `1;` `    ``long` `p = power(x, y / 2, M) % M;``    ``p = (p * p) % M;` `    ``return` `(y % 2 == 0) ? p : (x * p) % M;``  ``}`  `  ``// Function to find modular inverse``  ``// of a under modulo M``  ``// Assumption: M is prime``  ``static` `long` `modInverse(``long` `A, ``long` `M)``  ``{``    ``long` `g = gcd(A, M);``    ``return` `power(A, M - 2, M);``  ``}` `  ``// Function to Count number of ways``  ``// balls can be removed from following``  ``// boxes``  ``static` `long` `countWays(``long` `A, ``long` `K)``  ``{` `    ``// Precomputation array storing``    ``// factorials MOD 1e9 + 7``    ``long``[] fact=``new` `long``[100001];``    ``for``(``long` `i=0; i<100001; i++)``      ``fact[i]=0;` `    ``// Initialize the first element``    ``fact[1] = 1;`  `    ``// Filling factorial table``    ``for` `(``long` `i = 2; i <= 100000; i++) {``      ``fact[i] = (fact[i - 1] * i) % MOD;``    ``}` `    ``// Storing Total number of arrangements``    ``long` `ans = (fact[A + K] * modInverse(fact[A - 1], MOD))``      ``% MOD;` `    ``// Multiplying inverse of (K + 1)!``    ``ans = (ans * modInverse(fact[K + 1], MOD)) % MOD;` `    ``// Returning the answer``    ``return` `ans;``  ``}` `  ``// Driver Code``  ``static` `void` `Main(``string``[] args)``  ``{` `    ``// Input 1``    ``long` `A = 2, K = 2;` `    ``// Function Call``    ``Console.WriteLine(countWays(A, K));` `    ``// Input 2``    ``long` `A1 = 3, K1 = 3;` `    ``// Function Call``    ``Console.WriteLine(countWays(A1, K1));``  ``}``}` `// this code is contributed by poojaagarwal2.`

## Javascript

 `// JavaScript code to implement the approach` `// To avoid integer overflow``const MOD = BigInt(1e9 + 7);` `// To find GCD of a and b``function` `gcd(a, b) {``    ``if` `(a === 0n) {``        ``return` `b;``    ``}``    ``return` `gcd(b % a, a);``}` `// To compute x^y under modulo M``function` `power(x, y, M) {``    ``if` `(y === 0n) {``        ``return` `1n;``    ``}` `    ``var` `p = power(x, y / 2n, M) % M;``    ``p = (p * p) % M;` `    ``return` `(y % 2n === 0n) ? p : (x * p) % M;``}` `// Function to find modular inverse``// of a under modulo M``// Assumption: M is prime``function` `modInverse(A, M) {``    ``var` `g = gcd(A, M);``    ``return` `power(A, M - 2n, M);``}` `// Function to Count number of ways``// balls can be removed from following``// boxes``function` `countWays(A, K) {``    ``// Precomputation array storing``    ``// factorials MOD 1e9 + 7``    ``var` `fact = ``new` `Array(100001);``    ` `    ``// Initialize the first element``    ``fact[1] = 1n;``  ` `    ``// Filling factorial table``    ``for` `(``var` `i = 2n; i <= 100000n; i++) {``        ``fact[i] = (fact[i - 1n] * i) % MOD;``    ``}` `    ``// Storing Total number of arrangements``    ``var` `ans = (fact[A + K] * modInverse(fact[A - 1n], MOD)) % MOD;` `    ``// Multiplying inverse of (K + 1)!``    ``ans = (ans * modInverse(fact[K + 1n], MOD)) % MOD;` `    ``// Returning the answer``    ``return` `ans;``}` `// Driver code` `// Input 1``var` `A = 2n;``var` `K = 2n;` `// Function Call``console.log(countWays(A, K));` `// Input 2``var` `A1 = 3n;``var` `K1 = 3n;` `// Function Call``console.log(countWays(A1, K1));` `// This Code is Contributed by Prasad Kandekar(prasad264)`

Output

```4
15```

Time Complexity: O(N)
Auxiliary Space: O(N)

Related Articles:

My Personal Notes arrow_drop_up