Python code to solve the Car Talk puzzler about perfect square dance partners

Perfect Square Dance Solver

import math
import itertools

# This Python script solves the Car Talk puzzler about Perfect Square Dance
# http://www.cartalk.com/content/puzzler/transcripts/200842/index.html


def perfect_square(n):
  return math.sqrt(n) == int(math.sqrt(n))

def flatten(nested):
  """ Flatten one level of nesting.  Returns a generator object
  For instance:
  list(flatten([(1,3),(5,6)])) --> [1,3,5,6] """
  return itertools.chain.from_iterable(nested)

def all_guests_present_once(combination):
  """ Returns whether each guest is present once
  Combination is a list of tuples, e.g. [(1,5),(7,8)]
  """
  flattened = list(flatten(combination))
  return len(set(flattened)) == len(flattened)

def dance_arrangement(num_guests):
  """
  Returns a valid pairing for all guests if possible, else an empty set
  """
  # Clearly you need an even number of guests to have everyone paired
  if num_guests % 2 == 1:
    return []
  else:
    # range creates a list that's exclusive of last number.  Thus to go from 1 to n, 
    # use range(1, n+1)
    guests = range(1, num_guests + 1)
    # By enforcing the x<y constraint, I eliminate
    # a whole bunch of equivalent pairs (e.g. (3,6) and (6,3)).
    pairs = [(x,y) for x in guests for y in guests if x<y and perfect_square(x+y)]
    # brute force search
    all_arrangements = itertools.combinations(pairs, num_guests / 2)
    return filter(all_guests_present_once, all_arrangements)

def main():
  # Check all combinations from 2 to 20
  for i in range(2,21):
    valid_arrangements = dance_arrangement(i)
    if len(valid_arrangements) > 0:
      print i, valid_arrangements

if __name__ == '__main__':
  main()