python - Armstrong/Narcisstic number for base 16/hex

Question

What's the goal?

My goal is to find all armstrong/narcisstic numbers in hex for a given amount of digits.

The basic idea

The basic idea is that for a set of digits e.g. [A, 3, F, 5] the sum of powers is always the same no matter the order in which they occur. That means we don't have to look at every possible number up to our maximum which should greatly reduce runtime.

What I have so far

# Armstrong numbers base 16 for n digits

import time
import itertools
from typing import Counter

pows = [[]]

def genPow(max, base):
    global pows
    pows = [[0]*1 for i in range(base)]
    for i in range(base):
        pows[i][0] = i ** max

def check(a, b):
    c1 = Counter(a)
    c2 = Counter(b)

    diff1 = c1-c2
    diff2 = c2-c1

    # Check if elements in both 'sets' are equal in occurence
    return (diff1 == diff2)

def armstrong(digits):
    results = []
    genPow(digits, 16)
    # Generate all combinations without consideration of order
    for set in itertools.combinations_with_replacement('0123456789abcdef', digits):
        sum = 0
        # Genereate sum for every 'digit' in the set
        for digit in set:
            sum = sum + pows[int(digit, 16)][0] 
        # Convert to hex
        hexsum = format(sum, 'x')
        # No point in comparing if the length isn't the same
        if len(hexsum) == len(set):
            if check(hexsum, set):
                results.append(hexsum)

    return sorted(results)

start_time = time.time()
print(armstrong(10))
print("--- %s seconds ---" % (time.time() - start_time))

My problem

My issue is that this is still rather slow. It takes up to ~60 seconds for 10 digits. I'm pretty sure there are ways to do this more efficient. Some things I can think of, but don't know how to do are: faster way to generate combinations, condition for stopping calc. of sum, better way to compare the sum and set, convert to hex after comparing

Any ideas how to optimize this?

Edit: I tried to compare/check a bit differently and it's already a bit faster this way https://gist.github.com/Claypaenguin/d657c4413b510be580c1bbe3e7872624 Meanwhile I'm trying to understand the recursive approach, because it looks like it'll be a lot faster.

Answer 1

Your problem is that combinations_with_replacement for base b and length l is returning (b+l choose b) different things. Which in your case (base 16, length 10) means that you have 5,311,735 combinations.

Each of which you then do a heavyweight calculation on.

What you need to do is filter the combinations that you are creating as you are creating them. As soon as you realize that you are not on the way to an Armstrong number, abandon that path. The calculation will seem more complicated, but it is worthwhile when it lets you skip over whole blocks of combinations without having to individually generate them.

Here is pseudocode for the heart of the technique:

# recursive search for Armstrong numbers with:
#
#   base = base of desired number
#   length = length of desired number
#   known_digits = already chosen digits (not in order)
#   max_digit = the largest digit we are allowed to add
#
# The base case is that we are past or at a solution.
#
# The recursive cases are that we lower max_digit, or add max_digit to known_digits.
#
# When we add max_digit we compute min/max sums.  Looking at those we
# stop searching if our min_sum is too big or our max_sum is too small.
# We then look for leading digits in common.  This may let us discover
# more digits that we need.  (And if they are too big, we can't do that.)
def search(base, length, known_digits, max_digit):
    digits = known_digits.copy() # Be sure we do not modify the original.
    answer = []
    if length < len(digits):
        # We can't have any solutions.
        return []
    elif length == len(digits):
        if digits is a solution:
            return [digits]
        else:
            return []
    elif 0 < max_digit:
        answer = search(base, length, digits, max_digit-1)
    digits.append(max_digit)

    # We now have some answers, and known_digits.  Can we find more?
    find min_sum (all remaining digits are 0)
    if min_sum < base**(length-1):
        min_sum = base**(length-1)

    find max_sum (all remaining digits are max_digit)
    if base**length <= max_sum:
        max_sum = base**length - 1
     
    # Is there a possible answer between them?
    if max_sum < base**(length-1) or base**length <= min_sum:
        return answer # can't add more
    else:
        min_sum_digits = base_digits(min_sum, base)
        max_sum_digits = base_digits(max_sum, base)
        common_leading_digits = what digits are in common?
        new_digits = what digits in common_leading_digits can't be found in our known_digits?
        if 0 == len(new_digits):
            return answer + search(base, length, digits, max_digit)
        elif max_digit < max(new_digits):
            # Can't add this digit
            return answer
        else:
            digits.extend(new_digits)
            return answer + search(base, length, digits, max_digit)

---

I had a small logic error, but here is working code:

def in_base (n, b):
    answer = []
    while 0 < n:
        answer.append(n % b)
        n = n // b
    return answer


def powers (b, length, cached={}):
    if (b, length) not in cached:
        answer = []
        for i in range(b):
            answer.append(i**length)
        cached[(b, length)] = answer
    return cached[(b, length)]

def multiset_minus (a, b):
    count_a = {}
    for x in a:
        if x not in count_a:
            count_a[x] = 1
        else:
            count_a[x] += 1
    minus_b = []
    for x in b:
        if x in count_a:
            if 1 == count_a[x]:
                count_a.pop(x)
            else:
                count_a[x] -= 1
        else:
            minus_b.append(x)
    return minus_b


def armstrong_search (length, b, max_digit=None, known=None):
    if max_digit is None:
        max_digit = b-1
    elif max_digit < 0:
        return []

    if known is None:
        known = []
    else:
        known = known.copy() # Be sure not to accidentally share

    if len(known) == length:
        base_rep = in_base(sum([powers(b,length)[x] for x in known]), b)
        if 0 == len(multiset_minus(known, base_rep)):
            return [(base_rep)]
        else:
            return []
    elif length < len(known):
        return []
    else:
        min_sum = sum([powers(b,length)[x] for x in known])
        max_sum = min_sum + (length - len(known)) * powers(b,length)[max_digit]

        if min_sum < b**(length-1):
            min_sum = b**(length-1)
        elif b**length < min_sum:
            return []

        if b**length < max_sum:
            max_sum = b**length - 1
        elif max_sum < b**(length-1):
            return []

        min_sum_rep = in_base(min_sum, b)
        max_sum_rep = in_base(max_sum, b)

        common_digits = []
        for i in range(length-1, -1, -1):
            if min_sum_rep[i] == max_sum_rep[i]:
                common_digits.append(min_sum_rep[i])
            else:
                break

        new_digits = multiset_minus(known, common_digits)
        if 0 == len(new_digits):
            answers = armstrong_search(length, b, max_digit-1, known)
            known.append(max_digit)
            answers.extend(armstrong_search(length, b, max_digit, known))
            return answers
        else:
            known.extend(new_digits)
            return armstrong_search(length, b, max_digit, known)

And for a quick example:

digits = list('0123456789abcdef')
print([''.join(reversed([digits[i] for i in x])) for x in armstrong_search(10, len(digits))])

Takes a little over 2 seconds to find that the only answer is bcc6926afe.