Numbers of length N and value less than K I am solving interview question from here.

Problem : Given a set of digits (A) in sorted order, find how many numbers of length B are possible whose value is less than number C.
Constraints: 1 ≤ B ≤ 9, 0 ≤ C ≤ 1e9 &amp; 0 ≤ A[i] ≤ 9
Input: A = [ 0 1 5], B=  1  , C = 2 ; Output:  2 (0 and 1 are possible)  
Input: A = 0 1 2 5  , B =  2  , C = 21 ; Output: 5 (10, 11, 12, 15, 20 are possible)


This is my approach
from itertools import product
from itertools import ifilter
def solve(A, B, C):
    if A == [] or B &gt; len(str(C)):
        return 0
       
    elif B &lt; len(str(C)):
        #constraint is B
        if B == 1:
            new_list = A
            return len(new_list)
        else:
            new_list = list((product((''.join(str(i)for i in A)),repeat = B)))
            b = [''.join(num) for num in new_list]
            c = list(ifilter(lambda  x: x[0]!='0'  , b))
            return len(c)
                
            
    elif B == len(str(C)):
        #constraint is C 
        if B == 1:
            new_list = [i  for i in A if i&lt; C]
            return len(new_list)
        else:
            new_list = list((product((''.join(str(i)for i in A)),repeat = B)))
            b = [''.join(num) for num in new_list]
            c = list(ifilter(lambda  x: x[0]!='0' and int(x) &lt; C   , b))
            return len(c)

Test cases:
assert solve([2],5,51345) == 1
assert solve([],1,1) == 0
assert solve([ 2, 3, 5, 6, 7, 9 ],5,42950) == 2592
assert solve([0],1,5) == 1
assert solve([0,1,2,5],1,123) == 4
assert solve([0,1,5],1,2) == 2
assert solve([ 3 ],5, 26110) == 0
assert solve([0,1,2,5],2,21) == 5

How can I optimize this code in terms of memory usage?
Optimise memory usage

You could optimise memory usage by not converting your iterators into list and by avoiding non-required steps (like join).

Changing a few others details (formatting, adding tests, etc), you'd get something like:

from itertools import product
from itertools import ifilter

def solve(A, B, C):
    c_len = len(str(C))
    if A == [] or B &gt; c_len:
        return 0
    elif B &lt; c_len:
        # Constraint is B
        if B == 1:
            return len(A)
        else:
            candidates = product((str(i) for i in A), repeat = B)
            return sum(x[0] != '0' for x in candidates)
    else:
        assert B == c_len
        # Constraint is C
        if B == 1:
            return sum(i &lt; C for i in A)
        else:
            candidates = product((str(i) for i in A), repeat = B)
            return sum(x[0] != '0' and int(''.join(x)) &lt; C for x in candidates)

assert solve([2],5,51345) == 1
assert solve([],1,1) == 0
assert solve([2, 3, 5, 6, 7, 9],4,42950) == 1296
assert solve([2, 3, 5, 6, 7, 9],5,42950) == 2592
assert solve([0],1,5) == 1
assert solve([0,1,2,5],1,123) == 4
assert solve([0,1,5],1,2) == 2
assert solve([3],5, 26110) == 0
assert solve([0,1,2,5],2,21) == 5


Another algorithm

I'm pretty sure the whole thing can be optimised further by not generating the various numbers to count them but just using mathematical tricks to get the solution with no counting.

The easiest case to handle is B &lt; c_len:

elif B &lt; c_len:
    # All combinations of B elements are valid
    return len(set(A)) ** B


Actually, as mentionned by Maarten Fabré, this does not handle 0s perfectly. The code below is updated to handle it better.

The last case is trickier. We can try to use recursion to solve smaller versions of the problem. I didn't manage to make this work properly...

from itertools import product, ifilter, dropwhile, product, takewhile
import timeit

def solve_naive(A, B, C):
    A = set(str(A))
    mini = 10 ** (B-1)
    maxi = min(10 * mini, C)
    cand = [str(i) for i in (['0'] if B == 1 else []) + range(mini, maxi)]
    valid = [i for i in cand if all(c in A for c in i)]
    return len(valid)


def solve_op(A, B, C):
    # print(A, B, C)
    c_len = len(str(C))
    if A == [] or B &gt; c_len:
        return 0
    elif B &lt; c_len:
        # Constraint is B
        if B == 1:
            return len(A)
        else:
            candidates = product((str(i) for i in A), repeat = B)
            return sum(x[0] != '0' for x in candidates)
    else:
        assert B == c_len
        # Constraint is C
        if B == 1:
            return sum(i &lt; C for i in A)
        else:
            candidates = product((str(i) for i in A), repeat = B)
            return sum(x[0] != '0' and int(''.join(x)) &lt; C for x in candidates)


def solve_maarten(A, B, C):
    if A == [] or B &gt; len(str(C)):
        return 0
    c_tuple = tuple(map(int, str(C)))
    combinations = product(A, repeat=B)
    if B != 1:
        combinations = dropwhile(lambda x: x[0] == 0, combinations)
    if B == len(c_tuple):
        combinations = takewhile(lambda x: x &lt; c_tuple, combinations)
    combinations = list(combinations)
    return sum(1 for _ in combinations)


def solve(A, B, C):
    c_str = str(C)
    c_len = len(c_str)
    if A == [] or B &gt; c_len:
        return 0
    if B &lt; c_len:
        a_len = len(set(A))
        if B == 1:
            return a_len
        non_0_len = a_len - (0 in A)
        return non_0_len * (a_len ** (B-1))
    assert B == c_len    # Constraint is C
    head, tail = int(c_str[0]), c_str[1:]
    nb_first_dig_cand = sum(i &lt; head for i in A)
    if not tail or not nb_first_dig_cand:
        return nb_first_dig_cand
    if head in A:  # TODO: This case is not handled properly...
        # It should involve ret and solve(A, B-1, int(tail)) or something like that
        return solve_maarten(A, B, C)
    solve_c = solve(A, B-1, C)
    ret = nb_first_dig_cand * solve_c
    return ret


tests = [
    ([2], 4, 51345, 1),
    ([2], 5, 51345, 1),
    ([], 1, 1, 0),
    ([2, 3, 5, 6, 7, 9], 4, 42950, 1296),
    ([2, 3, 5, 6, 7, 9], 5, 42950, 2592),
    ([0], 1, 5, 1),
    ([0, 1, 2, 5], 1, 123, 4),
    ([0, 1, 5], 1, 2, 2),
    ([3], 5, 26110, 0),
    ([0, 1, 2, 5], 1, 21, 4),
    ([0, 1, 2, 5], 2, 21, 5),
    ([0, 1, 2, 5], 2, 201, 12),
    ([0, 1, 2, 5], 3, 2010, 48),
    ([0, 1, 2, 5], 4, 20108, 192),
    ([0, 1, 2, 5], 5, 201089, 768),
    ([0, 1, 2, 3, 4, 5, 7, 8], 5, 201089, 28672),
    ([0, 1, 2, 3, 4, 5, 7, 8], 6, 201089, 33344),
    ([0, 1, 2, 3, 4, 5, 7, 8, 9], 6, 200000, 59049),
    ([0, 1, 2, 3, 4, 5, 7, 8, 9], 6, 999999, 472391),
    ([1, 2, 3, 4, 5, 7, 8, 9], 6, 200000, 32768),
    ([1, 2, 3, 4, 5, 7, 8, 9], 6, 999999, 262143),
]
funcs = [solve, solve_op, solve_maarten, solve_naive]

for func in funcs:
    start = timeit.default_timer()
    for (A, B, C, exp) in tests:
        ret = func(A, B, C)
        if ret != exp:
            print "%s(%s, %d, %d): ret=%d, exp:%d" % (func.__name__, str(A), B, C, ret, exp)
    end = timeit.default_timer()
    print("Time for %s: %f" % (func.__name__, end - start))




def solve2(A, B, C):
    c_str = str(C)
    c_len = len(c_str)
    if A == [] or B &gt; c_len:
        return 0
    if B &lt; c_len:
        a_len = len(set(A))
        if B == 1:
            return a_len
        non_0_len = a_len - (0 in A)
        return non_0_len * (a_len ** (B-1))
    assert B == c_len    # Constraint is C
    head, last_dig = divmod(C, 10)
    nb_last_dig_cand = sum(i &lt; last_dig for i in A)
    if head == 0:
        return nb_last_dig_cand
    ret = solve_naive(A, B-1, head - 1) * len(A)
    ret_dummy = solve_naive(A, B, C)
    print(ret - ret_dummy, A, B, C)
    return ret_dummy