Quiz - discover X and Y given product and sum told to different people

gistfile1.py

"""
X and Y are two different integers, greater than 1, with sum less than 100. S and P are two mathematicians; S knows the sum X+Y, P knows the product X*Y, and both know the information in these two sentences. The following conversation occurs:

P says "I do not know X and Y."

S says "I knew you don't know X and Y."

P says "Now I know X and Y."

S says "Now I know X and Y too!"
"""

def sieve(n):
    "Return all primes <= n."
    np1 = n + 1
    s = list(range(np1)) # leave off `list()` in Python 2
    s[1] = 0
    sqrtn = int(round(n**0.5))
    for i in range(2, sqrtn + 1): # use `xrange()` in Python 2
        if s[i]:
            # next line:  use `xrange()` in Python 2
            s[i*i: np1: i] = [0] * len(range(i*i, np1, i))
    return filter(None, s)

PRIMES = sieve(100)

S_THINKS_P_MIGHT_KNOW = set( (x + y) for x in PRIMES for y in PRIMES if x<y and x+y < 100 )

REMAINING = set(range(5,100)) - S_THINKS_P_MIGHT_KNOW

def guess_xy_from_sum(s):
    return [(x, s-x) for x in range(2, s/2 + 1)]

"""
6 => (2,4) => X
11 => (2,9), (3,8), (4,7), (5,6)
       18     24     28     30
                            ~~
17 => (2,15), (3,14), (4,13), (5,12), (6,11), (7,10), (8,9)
       30      42      52      60      66      70      72
       ~~
...
Aha! Since P then knows X and Y, P can't be told a product like 30 which appears
at least twice (5+6 == 11, one of S's candidates, or 2+15==17)
restrict ourselves to those that appear only once
"""

p_candidates = {}
for s_candidate in REMAINING:
    xy_candidates = guess_xy_from_sum(s_candidate)
    for xy_candidate in xy_candidates:
        p_candidate = xy_candidate[0] * xy_candidate[1]
        other_xy_pairs = p_candidates.get(p_candidate, set())
        other_xy_pairs.add(xy_candidate)
        p_candidates[p_candidate] = other_xy_pairs
print "# of P candidates:", len(p_candidates)
    

# find p_candidates with only one possible sum
xy_candidates_unique = [list(p_candidates[p])[0]
                        for p in p_candidates
                        if len(p_candidates[p]) == 1]
print "# of XY candidates according to P:", len(xy_candidates_unique)

s_candidates = {}
# now do the flip side -- look at the sum and see which ones are unique
for xy_candidate in xy_candidates_unique:
    s_candidate = xy_candidate[0] + xy_candidate[1]
    other_xy_pairs = s_candidates.get(s_candidate, set())
    other_xy_pairs.add(xy_candidate)
    s_candidates[s_candidate] = other_xy_pairs
print "# of S candidates:", len(s_candidates)

# find s_candidates with only one possible product
xy_candidates_unique2 = [list(s_candidates[s])[0]
                         for s in s_candidates
                         if len(s_candidates[s]) == 1]
print "# of XY candidates according to S:", len(xy_candidates_unique2)
print "X=%d, Y=%d" % (xy_candidates_unique2[0][0], xy_candidates_unique2[0][1])

Answer: X=4, Y=13