Numbers whose squares can be rearranged to form exactly two other squares

shufflesquare.py

#!/usr/bin/env python
# 
# List the numbers n such that the digits of the square of n
# can be rearranged to form exactly two other squares. 
# For example, 13 belongs to list because 13*13 = 169, and 
# the digits of 169 can be rearranged to form the squares 196 
# (= 14*14) and 961 (= 31*31); but no other squares can be formed 
# by rearranging the digits of 169.
#
# Inspired by a question of James Tanton (@jamestanton on Twitter).
#
# Written by David Radcliffe (dradcliffe@gmail.com) on 21 April 2014.

MAX_DIGITS = 4 # Maximum number of digits in n

from collections import defaultdict
from time import time
start = time()

counter = defaultdict(list)

def sort_number(n):
    return ''.join(sorted(str(n)))
    
for n in xrange(1, 10**MAX_DIGITS):
    counter[sort_number(n*n)].append(n)
    
s = set()
for key, val in counter.iteritems():
    if len(val) == 3:
        s = s.union(set(val))

print sorted(s)
print "Time: %.2f seconds" % (time() - start)