Solution generator for the median-of-five problem

a_problem.txt

"Describe an algorithm that consumes five distinct comparable objects and
produces the median of these objects using a total of at most six comparisons."

Because creative thinking is hard, we use a simpler approach and try to compute
a suitable program recursively. Programs are represented by a tree of
comparisons at the internal nodes and resulting indices at the leaves.

The following code has been tested with Python 2.6.8, 2.7.3, and 3.3.0.

b_naive_generator.py

from itertools import combinations

def solve(n, m):
    """
    Determine a program to calculate the median of n distinct items using no
    more than m comparisons.

    n is assumed to be odd.
    """

    # Start with no known comparisons and all possible remaining
    # comparisons.
    return solve_helper(n, m, set(), set(combinations(range(n), 2)))

def solve_helper(n, m, cs, rem):
    """
    Determine a program to calculate the median of n distinct items using no
    more than m comparisons, given the set of true comparisons cs and the set
    of available comparisons rem.
    """

    # Find the median, if we can.
    for i in range(n):
        less, more = 0, 0

        for a, b in cs:
            if i == a:
                less += 1
            elif i == b:
                more += 1

        if less == n // 2 == more:
            return i

    # Not allowed any further comparisons.
    if m <= 0:
        return None

    # Attempt all the guesses possible at this point.
    for g in rem:
        rem_new = rem - set([g])

        t = solve_helper(n, m - 1, add_transitively(cs, g), rem_new)

        if t is None:
            continue

        f = solve_helper(n, m - 1, add_transitively(cs, (g[1], g[0])), rem_new)

        if f is None:
            continue

        return g, t, f

    return None

def add_transitively(cs, new):
    """
    Add the comparison new to the set of comparisons cs, along with any other
    comparisons that are implied by transitivity.
    """

    result = set(cs)
    queue = [new]

    while queue:
        n = queue.pop(0)
        result.add(n)

        for c in result:
            extra = None

            if n[0] == c[1]:
                extra = c[0], n[1]
            elif n[1] == c[0]:
                extra = n[0], c[1]

            if (extra is not None and
                    extra not in result and
                    extra not in queue):
                queue.append(extra)

    return result

if __name__ == '__main__':
    from pprint import pprint

    pprint(solve(5, 6))

c_check_solution.py

def run_program(p, i):
    """
    Run program p with input i.
    """

    try:
        # Is it an internal node?
        comparison = p[0]
    except TypeError:
        # No, it's a leaf.
        return i[p]

    # Continue down the correct subtree.
    if i[comparison[0]] < i[comparison[1]]:
        return run_program(p[1], i)
    else:
        return run_program(p[2], i)

if __name__ == '__main__':
    from itertools import permutations

    # This tree is produced in under a second on my machine by the naive
    # generator.
    p = ((0, 1),
         ((2, 3),
          ((1, 3),
           ((2, 4),
            ((1, 4), ((1, 2), 2, 1), ((0, 4), 4, 0)),
            ((1, 2), ((1, 4), 4, 1), ((0, 2), 2, 0))),
           ((0, 4),
            ((3, 4), ((0, 3), 3, 0), ((2, 4), 4, 2)),
            ((0, 3), ((0, 2), 2, 0), ((3, 4), 4, 3)))),
          ((1, 2),
           ((3, 4),
            ((1, 4), ((1, 3), 3, 1), ((0, 4), 4, 0)),
            ((1, 3), ((1, 4), 4, 1), ((0, 3), 3, 0))),
           ((0, 4),
            ((2, 4), ((0, 2), 2, 0), ((3, 4), 4, 3)),
            ((0, 2), ((0, 3), 3, 0), ((2, 4), 4, 2))))),
         ((2, 3),
          ((1, 2),
           ((0, 4),
            ((0, 2), ((2, 4), 2, 4), ((0, 3), 0, 3)),
            ((2, 4), ((3, 4), 3, 4), ((0, 2), 0, 2))),
           ((3, 4),
            ((1, 3), ((0, 3), 0, 3), ((1, 4), 1, 4)),
            ((1, 4), ((0, 4), 0, 4), ((1, 3), 1, 3)))),
          ((1, 3),
           ((0, 4),
            ((0, 3), ((3, 4), 3, 4), ((0, 2), 0, 2)),
            ((3, 4), ((2, 4), 2, 4), ((0, 3), 0, 3))),
           ((2, 4),
            ((1, 2), ((0, 2), 0, 2), ((1, 4), 1, 4)),
            ((1, 4), ((0, 4), 0, 4), ((1, 2), 1, 2))))))

    for i in permutations(range(5)):
        assert 2 == run_program(p, i), i

d_compact_solution.py

# Exploiting the symmetry of the produced tree (and with the help of some
# uglification tricks), it's possible to write a compact form of the solution
# in Python which "obviously" uses exactly 6 comparisons between the input
# items.

from operator import gt, lt

def median_of_five(L):
    a = [gt, lt][lt(L[0], L[1])]
    b = (6 - a(L[2], L[3])) // 2
    c, d = [(5 - b, b), (1, 0)][a(L[1], L[5 - b])]
    e, f = 4, abs(b - d)
    if a(L[abs(d - b)], L[4]): e, f = f, e
    g = a(L[f], L[c])
    h, i = [c, [e, f][g], d][g:g+2]
    return L[[h, i][a(L[h], L[i])]]

if __name__ == '__main__':
    from itertools import permutations

    for i in permutations(range(5)):
        assert 2 == median_of_five(i), i

    assert 'c' == median_of_five(['a', 'b', 'c', 'd', 'e'])
    assert 'c' == median_of_five(['b', 'e', 'a', 'd', 'c'])
    assert 'm' == median_of_five(['b', 'z', 'y', 'm', 'a'])

    assert (1, 2) == median_of_five([(1, 2), (3, 4), (0, 2), (1, 0), (5, 5)])