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 10⁹ + 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 = ^KC_K + ^K+1C_K + ^K+2C_K + ....... + ^K+A-1C_K
we know ^N+1C_M+1 = ^NC_M + ^NC_M+1 that is ^NC_M = ^N+1C_M+1 - ^NC_M+1
Total ways = ^K+1C_K+1 + (^K+2C_K+1 - ^K+1C_K+1) + (^K+3C_K+1 - ^K+2C_K+1) + ................ + (^K+AC_K+1 - ^K+A-1C_K-1)
Total ways = ^K+AC_K+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.

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:

Count ways choosing K balls from any given A boxes

Explore