Following a challenge on HackerRank, I want to find the sum of all numbers that have no more than x 4's, y 5's, and z 6's module 10^9 + 7. I'd like to know how I could improve this code below both for style and performance. Also, what's the canonical way to get unique permutations of a list that can include repeated elements?

import itertools
import sys
import collections

#create initial list based on input
x, y, z = map(int, sys.stdin.readline().split(' '))
permute_list = [4]*x + [5]*y + [6]*z

# collect iterators to permutations of all possible lengths
permuted = []
for i in range(1,(x+y+z+1)):
    permuted.append(itertools.permutations(permute_list, i))

# collect unique permutations from iterators
# there will be many repeats, anyway to avoid this inefficiency?
dynamic_table = collections.defaultdict(lambda: 0)
mod_var = (10**9 + 7)
for it in permuted:
    for i in it:
        new_sum = ''.join(map(str, i))
        if not dynamic_table[new_sum]:
            dynamic_table[new_sum] = int(new_sum)

print(sum(dynamic_table.values()) % mod_var)

I tried a few ways of doing this. What most surprised me is that implementing with a dictionary, as above, seems to beat a set implementation by about 30%, time-wise. Why is this? Also, what are some better ways to implement this and/or tweaks to the above code?


Imagine you had to do the same with 4 distinct digits, a, b, c and d, and lets forget for now about numbers with less than 4 digits. There are 4! arrangements of these 4 digits, of which 3! each will have the a in the first, second, third and fourth positions. So the contribution of the digit a to the overall sum will be a * 3! * 1111, and the total sum of all numbers generated from these four digits ends up being (a + b + c + d) * 3! * 1111. It is easy to check that this works comparing to a brute-force solution:

from itertools import  permutations

def sum_them_quick(a, b, c, d):
    return (a + b + c + d) * 3 * 2 * 1 * 1111

def sum_them_slow(a, b, c, d):
    total = 0
    for digits in permutations([a, b, c, d], 4):
        total += int(''.join(str(d) for d in digits))
    return total

If you test this, you will find that, while they both produce the same result, the performance is already 100x faster for this small example:

>>> sum_them_quick(4, 5, 6, 7) ==  sum_them_slow(4, 5, 6, 7)
>>> %timeit sum_them_quick(4, 5, 6, 7)
1000000 loops, best of 3: 349 ns per loop
>>> %timeit sum_them_slow(4, 5, 6, 7)
10000 loops, best of 3: 54.9 µs per loop

Your case is not that simple, because you have to handle the repeats, but it isn't that much harder either. If you have x 4's, y 5's and z 6's, there are (x + y + z)! / x! / y! / z! different arrangements of those digits to form unique numbers. And similarly to the previous case, a 4 will show up in any given position x / (x + y + z) of the times. So similarly, the following function will compute the total sum in constant time:

from math import factorial    

def sum_them_456(x, y, z):
    total = 0
    for j in range(x + y + z):
        total *= 10
        total += 1
    total *= 4*x + 5*y + 6*z
    total *= factorial(x + y + z - 1)
    total //= factorial(x) * factorial(y) * factorial(z)
    return total

There are cleverer ways of computing a quotient of factorials, but that is left as exercise.

You also have to handle all the possible permutations with less than the full x + y + z possible digits. There may be a better way of not having to brute-force your way through this, but I'm out of cleverness right now, so I am going to go with something very unsophisticated:

def sum_them_all_456(x, y, z):
    n = x + y + z
    total = 0
    trials = set()
    trials.add((x, y, z))
    while n > 0:
        new_trials = set()
        for trial in trials:
            total += sum_them_456(*trial)
            for idx in range(len(trial)):
                if trial[idx] > 0:
                    new_trial = trial[:idx] + (trial[idx] - 1,) + trial[idx+1:]
        trials = new_trials
        n -= 1
    return total

It still manages to handle values of x, y or z in the tenths in the blink of an eye, but there has to be a better way for this.

  • \$\begingroup\$ Awesome, that first part is very nice. And then for quotients of factorials, I guess I'd use numpy.gammaprod(). It's very embarassing that I did not even think about using combinatorics rather than brute-forcing \$\endgroup\$ – sunny Jul 17 '15 at 20:53
  1. Why sys.stdin... instead of input
  2. Why build the list of permutations instead of just iterating over it?
  3. Why a defaultdict instead of just a dict?

It's hard to tell why the set implementation is slower without seeing it.

  • \$\begingroup\$ (1) stdin was required by the problem. How is input more advantageous? (2) That's a good point. I will try to do this in one 'loop'. (3) I used a defaultdict so as not to worry about a key-value error and having to check that. Is it so bad? \$\endgroup\$ – sunny Jul 17 '15 at 20:45
  • \$\begingroup\$ (1) input is shorter and the more common way of doing this. (3) Yes. Letting something silently fail is always bad. If something unexpected happens, you want to know about it rather than letting your code continue. If you disagree, put a blanket try/except around all your code and try debugging it. \$\endgroup\$ – raylu Jul 17 '15 at 20:52
  • \$\begingroup\$ for (3) is there a performance difference? Here I expect key-value errors as part of the design, so it seemed sensible to start with a defaultdict that would simplify the checking compared to if...in dict.keys(). To me the latter is uglier, but if it's a performance issue I see your point. \$\endgroup\$ – sunny Jul 17 '15 at 20:56
  • \$\begingroup\$ You explicitly check if not dynamic_table[new_sum]:, which creates a (new_sum, 0) pair inside the dict. The normal way is to just if new_sum not in dynamic_table (no need for .keys()). The performance of this constant-time operation is pretty irrelevant. \$\endgroup\$ – raylu Jul 18 '15 at 22:02

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.