lock combinations

comb.py

from collections import Counter

xs = [i for i in range(5) for k in range(5)]


def next_z(ys, zs):
    cur_len = len(zs)
    if cur_len == 25:
        return

    yz_pairs = set(zip(ys, zs))
    xz_pairs = set(zip(xs, zs))
    next_x = xs[cur_len]
    next_y = ys[cur_len]
    nums = Counter(zs)

    avail_z = [z for z in range(5) if nums[z] < 5 and
               (next_x, z) not in xz_pairs and
               (next_y, z) not in yz_pairs]

    for z in avail_z:
        yield z


def next_y(ys):
    cur_len = len(ys)
    if cur_len == 25:
        return

    xy_pairs = set(zip(xs, ys))
    next_x = xs[cur_len]
    nums = Counter(ys)

    avail_y = [y for y in range(5) if nums[y] < 5 and
               (next_x, y) not in xy_pairs]

    for y in avail_y:
        yield y


def arrange_ys(ys):
    if len(ys) == 25:
        yield ys

    for y in next_y(ys):
        for solution_y in arrange_ys(ys + [y]):
            yield solution_y


def arrange_zs(ys, zs):
    if len(zs) == 25:
        yield list(zip(xs, ys, zs))

    for z in next_z(ys, zs):
        for solution in arrange_zs(ys, zs + [z]):
            yield solution


def find_queen():
    for ys in arrange_ys([]):
        n = 0
        print('ys:', ys)
        for solution in arrange_zs(ys, []):
            n += 1
            if n // 5000 * 5000 == n:
                print(n)
            yield solution
        print('total:', n)


def check(qs, n):
    tried = set()
    for x, y, z in qs:
        for i in range(n):
            tried.add((i, y, z))
            tried.add((x, i, z))
            tried.add((x, y, i))
    if len(tried) == n ** 3:
        return True


if __name__ == '__main__':
    for solution in find_queen():
        solution_10 = solution + [(9 - x, 9 - y, 9 - z)
                                  for x, y, z in solution]
        # print(','.join("%s%s%s" % (x, y, z) for x, y, z in solution_10))
        # if not check(solution_10, 10):
        #     print('.......')