Numbers of Length N having digits A and B and whose sum of digits contain only digits A and B | Python 3 implementation of the approach ; Function that returns true if the num contains a and b digits only ; Modular Exponentiation ; Function to return the modular inverse of x modulo MOD ; Function to return the required count of numbers ; Generating factorials of all numbers ; Generating inverse of factorials modulo MOD of all numbers ; Keeping a as largest number ; Iterate over all possible values of s and if it is a valid S then proceed further ; Check for invalid cases in the equation ; Find answer using combinatorics ; Add this result to final answer ; Driver Code
MAX = 100005 ; NEW_LINE MOD = 1000000007 NEW_LINE def check ( num , a , b ) : NEW_LINE INDENT while ( num ) : NEW_LINE INDENT rem = num % 10 NEW_LINE num = int ( num / 10 ) NEW_LINE if ( rem != a and rem != b ) : NEW_LINE INDENT return 0 NEW_LINE DEDENT DEDENT return 1 NEW_LINE DEDENT def power ( x , y ) : NEW_LINE INDENT ans = 1 NEW_LINE while ( y ) : NEW_LINE INDENT if ( y & 1 ) : NEW_LINE INDENT ans = ( ans * x ) % MOD NEW_LINE DEDENT y >>= 1 NEW_LINE x = ( x * x ) % MOD NEW_LINE DEDENT return ans % MOD NEW_LINE DEDENT def modInverse ( x ) : NEW_LINE INDENT return power ( x , MOD - 2 ) NEW_LINE DEDENT def countNumbers ( n , a , b ) : NEW_LINE INDENT fact = [ 0 for i in range ( MAX ) ] NEW_LINE inv = [ 0 for i in range ( MAX ) ] NEW_LINE ans = 0 NEW_LINE fact [ 0 ] = 1 NEW_LINE for i in range ( 1 , MAX , 1 ) : NEW_LINE INDENT fact [ i ] = ( 1 * fact [ i - 1 ] * i ) NEW_LINE fact [ i ] %= MOD NEW_LINE DEDENT inv [ MAX - 1 ] = modInverse ( fact [ MAX - 1 ] ) NEW_LINE i = MAX - 2 NEW_LINE while ( i >= 0 ) : NEW_LINE INDENT inv [ i ] = ( inv [ i + 1 ] * ( i + 1 ) ) NEW_LINE inv [ i ] %= MOD NEW_LINE i -= 1 NEW_LINE DEDENT if ( a < b ) : NEW_LINE INDENT temp = a NEW_LINE a = b NEW_LINE b = temp NEW_LINE DEDENT for s in range ( n , 9 * n + 1 , 1 ) : NEW_LINE INDENT if ( check ( s , a , b ) == 0 ) : NEW_LINE INDENT continue NEW_LINE DEDENT if ( s < n * b or ( s - n * b ) % ( a - b ) != 0 ) : NEW_LINE INDENT continue NEW_LINE DEDENT numDig = int ( ( s - n * b ) / ( a - b ) ) NEW_LINE if ( numDig > n ) : NEW_LINE INDENT continue NEW_LINE DEDENT curr = fact [ n ] NEW_LINE curr = ( curr * inv [ numDig ] ) % MOD NEW_LINE curr = ( curr * inv [ n - numDig ] ) % MOD NEW_LINE ans = ( ans + curr ) % MOD NEW_LINE DEDENT return ans NEW_LINE DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT n = 3 NEW_LINE a = 1 NEW_LINE b = 3 NEW_LINE print ( countNumbers ( n , a , b ) ) NEW_LINE DEDENT