function that counts...

Richard Thomas chardster at gmail.com
Wed May 19 21:36:22 EDT 2010


For this kind of problem you should avoid all that stringification. I
find it best to deal with sequences of digits of a fixed length and go
from there. For example:

def count1(m, n, cache={}):
    """Number of digit sequences of length `n` summing to `m`."""
    if n < 0 or m < 0:
        return 0
    elif n == 0:
        return int(m == 0)
    elif (m, n) in cache:
        return cache[m, n]
    # This is an optimisation using the combinatoric choose function.
    #elif m < 10:
    #    result = choose(n + m - 1, n - 1)
    else:
        result = 0
        for digit in xrange(min(10, m + 1)):
            result += count1(m - digit, n - 1)
    cache[m, n] = result
    return result

Notice the caching of results. With this we can compute the required
thing quite easily:

def count2(m, n):
    """Number of numbers less than `n` whose digits sum to `m`."""
    result = 0
    digits = map(int, str(n))
    length = len(digits)
    for idx, digit in enumerate(digits):
        for seq_digit in xrange(digit):
            seq_limit = m - seq_digit
            seq_length = length - idx - 1
            result += count1(seq_limit, seq_length)
        m -= digit
    return result

Essentially we move through the number left to right, choose digits
less than the digit in the number and count digit sequences to fill
the remainder. Then fix the actual digit and step forward.

An approach like this avoids the relatively slow stringification
process and makes decent use of caching and iteration (both easy and
efficient in Python).

In [1]: count2(25, 10000)
Out[1]: 348

In [2]: timeit count2(25, 10000)
100000 loops, best of 3: 12.6 us per loop



More information about the Python-list mailing list