1# Source: https://github.com/python/pyperformance
2# License: MIT
3
4# Simple, brute-force N-Queens solver.
5# author: collinwinter@google.com (Collin Winter)
6# n_queens function: Copyright 2009 Raymond Hettinger
7
8# Pure-Python implementation of itertools.permutations().
9def permutations(iterable, r=None):
10    """permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)"""
11    pool = tuple(iterable)
12    n = len(pool)
13    if r is None:
14        r = n
15    indices = list(range(n))
16    cycles = list(range(n - r + 1, n + 1))[::-1]
17    yield tuple(pool[i] for i in indices[:r])
18    while n:
19        for i in reversed(range(r)):
20            cycles[i] -= 1
21            if cycles[i] == 0:
22                indices[i:] = indices[i + 1 :] + indices[i : i + 1]
23                cycles[i] = n - i
24            else:
25                j = cycles[i]
26                indices[i], indices[-j] = indices[-j], indices[i]
27                yield tuple(pool[i] for i in indices[:r])
28                break
29        else:
30            return
31
32
33# From http://code.activestate.com/recipes/576647/
34def n_queens(queen_count):
35    """N-Queens solver.
36    Args: queen_count: the number of queens to solve for, same as board size.
37    Yields: Solutions to the problem, each yielded value is a N-tuple.
38    """
39    cols = range(queen_count)
40    for vec in permutations(cols):
41        if queen_count == len(set(vec[i] + i for i in cols)) == len(set(vec[i] - i for i in cols)):
42            yield vec
43
44
45###########################################################################
46# Benchmark interface
47
48bm_params = {
49    (50, 25): (1, 5),
50    (100, 25): (1, 6),
51    (1000, 100): (1, 7),
52    (5000, 100): (1, 8),
53}
54
55
56def bm_setup(params):
57    res = None
58
59    def run():
60        nonlocal res
61        for _ in range(params[0]):
62            res = len(list(n_queens(params[1])))
63
64    def result():
65        return params[0] * 10 ** (params[1] - 3), res
66
67    return run, result
68