I'm going through the exercises in Coding the Matrix, and have gotten to this task: write a procedure all_3_digit_numbers(base, digits), where input is a positive integer base and digits is the set {0, 1,2, ..., base-1}, and the output is the set of all three-digit numbers where the base is base. I was able to come up with one solution:
def all_3_digit_numbers(base, digits): return {a+b+c for a in [d*base**2 for d in digits] for b in [d*base**1 for d in digits] for c in [d*base**0 for d in digits]}
which fits in one line. But I thought the repetition of the expression d*base**exp could be consolidated. So I also have this solution:
def all_3_digit_numbers(base, digits):
[alist,blist,clist] = [[digit*base**exp for digit in digits] for exp in [0,1,2]]
return {a+b+c for a in alist for b in blist for c in clist}
which is more elegant (I guess), but requires an extra expression. I was wondering if there's a way to get the best of both worlds - a comprehension that doesn't require variable assignment in a separate statement as it does the second solution, but where the list of digits to add is not repeated as it is in the first solution.
These intermediate lists make the code really long. No need for that:
def all_3_digit_numbers(base, digits):
return {a*base**2 + b*base + c for a in digits for b in digits for c in digits}
And if you want to be a little bit more fancy, use itertools:
from itertools import product
def all_3_digit_numbers2(base, digits):
return {a*base**2 + b*base + c for a, b, c in product(digits, repeat=3)}