balance_puzzle_13.py

def compare(a, b):
    return (a > b) - (a < b)

weights = {i : 10 for i in range(13)}

def weigh(left, right):
    assert len(left) == len(right)
    return compare(sum(weights[i] for i in left), sum(weights[i] for i in right))

def solve_2_less_1r(left, right):
    assert len(left) == 2
    assert len(right) == 1

    v1, v3 = list(left)
    v2 = list(right)[0]

    c = weigh([v1], [v3])
    if c == 0:
        return v2
    elif c == -1:
        return v1
    else:
        return v3

def solve_2_less(a, b, r):
    c = weigh([a], [r])
    if c == 0:
        return b
    else:
        return a

def solve_2_gt_1r(left, right):
    return {
        left[0] : left[1],
        left[1] : left[0],
        right[0] : right[0],
    }[solve_2_less_1r(left, right)]

def solve_4_less_4(left4, right4, r):
    left = [left4[0], right4[1], r]
    right = [right4[0], left4[1], right4[2]]
    c = weigh(left, right)
    if c == 0:
        return solve_2_less_1r([left4[2], left4[3]], [right4[3]])
    elif c == -1:
        return solve_2_gt_1r([right4[0], right4[2]], [left4[0]])
    else:
        return solve_2_less(left4[1], right4[1], r)


def solve():
    vals = list(weights.keys())
    left = vals[:4]
    right = vals[4:8]

    c = weigh(left, right)

    if c == 0:
        left = vals[8:10]
        right = [vals[10], vals[0]]
        c = weigh(left, right)
        if c == 0:
            return solve_2_less(vals[11], vals[12], vals[0])
        elif c == -1:
            return solve_2_less_1r(left, [vals[10]])
        else:
            return solve_2_gt_1r(left, [vals[10]])
    elif c == -1:
        return solve_4_less_4(left, right, vals[-1])
    else:
        return solve_4_less_4(right, left, vals[-1])

for i in weights.keys():
    weights[i] = 9
    ok_less = {i : 'ok', None: 'n/a'}.get(solve(), 'INCORRECT')
    weights[i] = 11
    ok_greater = {i : 'ok', None: 'n/a'}.get(solve(), 'INCORRECT')
    print(i, ok_less, ok_greater)
    weights[i] = 10