two_students_riddle.py

#!/usr/bin/env python3
"""Problem statement:

There is a teacher and 2 students in a classroom. The students are A and B.
The teacher thinks of 2 positive integers and tells the sum of those numbers
to student A without student B hearing it. Then tells their product to student
B without student A hearing it. After this, the teacher asks the 2 students
what was the 2 numbers.

First student A says: I don't know.
Then student B says: I don't know either.
After hearing this, student A says: Now I know.
Then student B says: Now I know them too.

What were the 2 numbers?
"""


def memoize(f):
    memo = {}
    def helper(x):
        if x not in memo:            
            memo[x] = f(x)
        return memo[x]
    return helper


def _tup(x, y):
    return (x, y) if x < y else (y, x)


@memoize
def add_decomp(s):
    return {_tup(i, s - i) for i in range(1, s)}


@memoize
def prod_decomp(p):
    return {_tup(i, p // i) for i in range(1, p) if p % i == 0}


@memoize
def a_first(s):
    return len(add_decomp(s)) > 1


@memoize
def b_first(p):
    return sum([a_first(i + j) for i, j in prod_decomp(p)]) > 1


@memoize
def a_second(s):
    return sum([b_first(i * j) for i, j in add_decomp(s)]) == 1


@memoize
def b_second(p):
    return sum([a_second(i + j) for i, j in prod_decomp(p)]) == 1


log = False
maxv = 10


for a in range(1, maxv):
    for b in range(1, a + 1):
        s, p = a + b, a * b
        if log:
            print(f'a: {a}, b: {b}, s: {s}, p: {p}', end=' :: ')
        if not a_first(p):
            if log:
                print('a knows the first time')
            continue
        if not b_first(s):
            if log:
                print('b knows the first time')
            continue
        if not a_second(p):
            if log:
                print('a does not know the second time')
            continue
        if not b_second(s):
            if log:
                print('b does not know the second time')
            continue
        print(f'a: {a}, b: {b} works!')