Open In App

Count ways choosing K balls from any given A boxes

Improve
Improve
Like Article
Like
Save
Share
Report

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.
  • Print the answer.

Below is the implementation of the above approach:

C++




// C++ code to implement the approach
#include <bits/stdc++.h>
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:



Last Updated : 18 Feb, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads