Counting numbers with given digits and digit sum

Given a number N, count the numbers X of length exactly N such that the number X and the sum of digits of the number X have digits A and B only in their decimal representation. The length of a number is defined as the number of digits in its decimal representation without leading zeroes.

Note: As this count may be large, return the remainder after dividing it by 1000000007 (10⁹ + 7).

Example:

Input: A = 1, B = 2, N = 2
Output: 1
Explanation: The numbers of length 2 having only digits 1 and 2 are: 11, 12, 21, 22 out of which only 11 has the sum of digits = 2.

Input: A = 2, B = 4, N = 6
Output: 7
Explanation: The numbers of length 2 having only digits 2 and 4 are: 244444, 424444, 442444, 444244, 444424, 444442, 444444 and their sum of digits are: 22, 22, 22, 22, 22, 22, 24 respectively.

Approach: To solve the problem, follow the below idea:

The idea is to take only those numbers which have digits A and B and then determine if their sum also has digits A and B only. For a number of length N every i occurrence of A, will have (N – i) occurrences of B. If the sum (A * i + B * (N – i)) forms a valid number, then we can calculate all possible combinations of A and B using the formula N! / ((N − i)! * i!)​ , which represents the binomial coefficient. We precompute the factorial and inverse factorial of all numbers up to N in advance. The inverse factorial is used to calculate the denominator of the binomial coefficient in a modular arithmetic. In modular arithmetic, we can’t simply divide by a number. Instead, we multiply by its modular multiplicative inverse. The function power(fact[i], MOD – 2), using Fermat Little Theorem, calculates this inverse using Binary Exponentiation. The count of such numbers, given by the binomial coefficient, is then added to the final answer.

Step-by-step algorithm:

• Maintain a Function power(a, b) to calculate a^b under modulo MOD using binary exponentiation.
• Generate all numbers in the order of occurrences of digit A:
  • 0 occurrence of digit A, N occurrences of digit B.
  • 1 occurrence of digit A, (N – 1) occurrences of digit B.
  • 2 occurrence of digit A, (N – 2) occurrence of digit B and so on.
• Check for the sum of digits of all the above numbers and if the sum has only digits A and B, count all the numbers that can be formed by rearranging all the digits.
• Sum of all the above numbers will be the final answer.

Below is the implementation of the algorithm:

C++

// C++ Code
#include <bits/stdc++.h>
using namespace std;
#define MOD 1000000007

// Function to calculate a^b under modulo MOD using 
// binary exponentiation
long long power(long long a, long long b)
{
    // If b = 0, whatever be the value of a, 
    // our result will be 1.
    long long res = 1;
    while (b > 0) {
        // If b is an odd number, then 
        // (a^b) = (a * (a^(b–1)/2)^2)
        if (b & 1) {
            res = (res * a) % MOD;
        }

        // If b is an even number, then 
        // (a^b) = ((a^2)^(b/2))
        a = (a * a) % MOD;
        b >>= 1;
    }
    return res;
}

// Function to check if a number is good
bool check(long long N, long long A, long long B)
{
    while (N) {
        // If any digit is neither A nor B, the number is
        // not valid
        if (N % 10 != A && N % 10 != B)
            return 0;

        N /= 10;
    }
    // All digits are either A or B
    return 1; 
}

long long excellentno(long long A, long long B, long long N)
{
    vector<long long> fact(N + 1), inv(N + 1);
    fact[0] = inv[0] = 1;

    // Precompute factorials and their inverses
    for (long long i = 1; i <= N; i++) {
        // Compute factorial
        fact[i] = (fact[i - 1] * i) % MOD;

        // Compute inverse factorial
        inv[i] = power(fact[i], MOD - 2);
    }

    long long ans = 0;
    for (long long i = 0; i <= N; i++) {
        // If the sum of digits is a good number, add the
        // count to the answer
        if (check(A * i + B * (N - i), A, B)) {
            // Total combinations = N! / (i!(N-i)!) = N! *
            // inv[i] * inv[N-i]
            ans = (ans
                + (fact[N] * inv[i] % MOD * inv[N - i]
                    % MOD))
                % MOD;
        }
    }
    return ans;
}

int main()
{
    long long A, B, N;
    // Example 1
    A = 1, B = 2, N = 2;
    cout << excellentno(A, B, N) << endl;

    // Example 2
    A = 2, B = 4, N = 6;
    cout << excellentno(A, B, N) << endl;
    return 0;
}

Java

import java.util.*;

public class ExcellentNumber {
    static final long MOD = 1000000007;

    // Function to calculate a^b under modulo MOD using
    // binary exponentiation
    static long power(long a, long b)
    {
        // If b = 0, whatever be the value of a, our result
        // will be 1.
        long res = 1;
        while (b > 0) {
            // If b is an odd number, then (a^b) = (a *
            // (a^(b–1)/2)^2)
            if ((b & 1) == 1) {
                res = (res * a) % MOD;
            }

            // If b is an even number, then (a^b) =
            // ((a^2)^(b/2))
            a = (a * a) % MOD;
            b >>= 1;
        }
        return res;
    }

    // Function to check if a number is good
    static boolean check(long N, long A, long B)
    {
        while (N > 0) {
            // If any digit is neither A nor B, the number
            // is not valid
            if (N % 10 != A && N % 10 != B)
                return false;

            N /= 10;
        }
        // All digits are either A or B
        return true;
    }

    static long excellentNo(long A, long B, long N)
    {
        long[] fact = new long[(int)(N + 1)];
        long[] inv = new long[(int)(N + 1)];
        fact[0] = inv[0] = 1;

        // Precompute factorials and their inverses
        for (long i = 1; i <= N; i++) {
            // Compute factorial
            fact[(int)i] = (fact[(int)(i - 1)] * i) % MOD;

            // Compute inverse factorial
            inv[(int)i] = power(fact[(int)i], MOD - 2);
        }

        long ans = 0;
        for (long i = 0; i <= N; i++) {
            // If the sum of digits is a good number, add
            // the count to the answer
            if (check(A * i + B * (N - i), A, B)) {
                // Total combinations = N! / (i!(N-i)!) = N!
                // * inv[i] * inv[N-i]
                ans = (ans
                       + (fact[(int)N] * inv[(int)i] % MOD
                          * inv[(int)(N - i)] % MOD))
                      % MOD;
            }
        }
        return ans;
    }

    public static void main(String[] args)
    {
        long A, B, N;

        // Example 1
        A = 1;
        B = 2;
        N = 2;
        System.out.println(excellentNo(A, B, N));

        // Example 2
        A = 2;
        B = 4;
        N = 6;
        System.out.println(excellentNo(A, B, N));
    }
}

// This code is contributed by akshitaguprzj3

C#

using System;
using System.Numerics;

public class Program
{
    const long MOD = 1000000007;

    // Function to calculate a^b under modulo MOD using 
    // binary exponentiation
    static long Power(long a, long b)
    {
        // If b = 0, whatever be the value of a, 
        // our result will be 1.
        long res = 1;
        while (b > 0)
        {
            // If b is an odd number, then 
            // (a^b) = (a * (a^(b–1)/2)^2)
            if ((b & 1) == 1)
            {
                res = (res * a) % MOD;
            }

            // If b is an even number, then 
            // (a^b) = ((a^2)^(b/2))
            a = (a * a) % MOD;
            b >>= 1;
        }
        return res;
    }

    // Function to check if a number is good
    static bool Check(long N, long A, long B)
    {
        while (N > 0)
        {
            // If any digit is neither A nor B, the number is
            // not valid
            if (N % 10 != A && N % 10 != B)
                return false;

            N /= 10;
        }
        // All digits are either A or B
        return true;
    }

    static long ExcellentNo(long A, long B, long N)
    {
        long[] fact = new long[N + 1];
        long[] inv = new long[N + 1];
        fact[0] = inv[0] = 1;

        // Precompute factorials and their inverses
        for (long i = 1; i <= N; i++)
        {
            // Compute factorial
            fact[i] = (fact[i - 1] * i) % MOD;

            // Compute inverse factorial
            inv[i] = Power(fact[i], MOD - 2);
        }

        long ans = 0;
        for (long i = 0; i <= N; i++)
        {
            // If the sum of digits is a good number, add the
            // count to the answer
            if (Check(A * i + B * (N - i), A, B))
            {
                // Total combinations = N! / (i!(N-i)!) = N! *
                // inv[i] * inv[N-i]
                ans = (ans + (fact[N] * inv[i] % MOD * inv[N - i] % MOD)) % MOD;
            }
        }
        return ans;
    }

    public static void Main(string[] args)
    {
        long A, B, N;
        // Example 1
        A = 1; B = 2; N = 2;
        Console.WriteLine(ExcellentNo(A, B, N));

        // Example 2
        A = 2; B = 4; N = 6;
        Console.WriteLine(ExcellentNo(A, B, N));
    }
}

JavaScript

// Define the modulo constant
const MOD = BigInt(1000000007);

// Function to calculate a^b under modulo MOD using binary exponentiation
function power(a, b) {
    let res = BigInt(1);
    a = BigInt(a);
    b = BigInt(b);
    while (b > 0) {
        if (b & 1n) {
            res = (res * a) % MOD;
        }
        a = (a * a) % MOD;
        b >>= 1n;
    }
    return res;
}

// Function to check if a number is good
function check(N, A, B) {
    while (N) {
        if (N % 10n != BigInt(A) && N % 10n != BigInt(B))
            return false;
        N /= 10n;
    }
    return true; 
}

function excellentno(A, B, N) {
    let fact = new Array(N + 1).fill(BigInt(1));
    let inv = new Array(N + 1).fill(BigInt(1));

    for (let i = 1; i <= N; i++) {
        fact[i] = (fact[i - 1] * BigInt(i)) % MOD;
        inv[i] = power(fact[i], MOD - 2n);
    }

    let ans = BigInt(0);
    for (let i = 0; i <= N; i++) {
        if (check(BigInt(A) * BigInt(i) + BigInt(B) * BigInt(N - i), A, B)) {
            ans = (ans + (fact[N] * inv[i] % MOD * inv[N - i] % MOD)) % MOD;
        }
    }
    return ans.toString();
}

let A = 1, B = 2, N = 2;
console.log(excellentno(A, B, N));

A = 2, B = 4, N = 6;
console.log(excellentno(A, B, N));

Python3

MOD = 1000000007

# Function to calculate a^b under modulo MOD using binary exponentiation
def power(a, b):
    # If b = 0, whatever be the value of a, our result will be 1.
    res = 1
    while b > 0:
        # If b is an odd number, then (a^b) = (a * (a^(b–1)/2)^2)
        if b & 1:
            res = (res * a) % MOD
        # If b is an even number, then (a^b) = ((a^2)^(b/2))
        a = (a * a) % MOD
        b >>= 1
    return res

# Function to check if a number is good
def check(N, A, B):
    while N:
        # If any digit is neither A nor B, the number is not valid
        if N % 10 != A and N % 10 != B:
            return False
        N //= 10
    # All digits are either A or B
    return True

def excellentno(A, B, N):
    fact = [0] * (N + 1)
    inv = [0] * (N + 1)
    fact[0] = inv[0] = 1

    # Precompute factorials and their inverses
    for i in range(1, N + 1):
        # Compute factorial
        fact[i] = (fact[i - 1] * i) % MOD

        # Compute inverse factorial
        inv[i] = power(fact[i], MOD - 2)

    ans = 0
    for i in range(N + 1):
        # If the sum of digits is a good number, add the count to the answer
        if check(A * i + B * (N - i), A, B):
            # Total combinations = N! / (i!(N-i)!) = N! * inv[i] * inv[N-i]
            ans = (ans + (fact[N] * inv[i] % MOD * inv[N - i] % MOD)) % MOD

    return ans

# Example 1
A, B, N = 1, 2, 2
print(excellentno(A, B, N))

# Example 2
A, B, N = 2, 4, 6
print(excellentno(A, B, N))

Output

1
7

Time Complexity: O(N * logN), where N is the length of numbers.
Auxiliary Space: O(N)