Functional n-queens

Question

Following the book "Structure and interpretation of computer programs" I have tried to implement a functional solution to the problem of N-queens (implemented by the function nQueens). However, I am not satisfied by the reliance on for loops and the general readability of the code. Could I get your opinion the code?

import functools as fn

def safe(positions, k):
    """ Given a list of positions, returns if the queen on the k column
    is safe. """
    enumeratedPos = tuple(enumerate(positions));
    menacesK = fn.partial(menaces, enumeratedPos[k]);
    return not any(map(menacesK,
        filter(lambda pair: pair[0] != k, enumeratedPos)));

def addNewColumn(positions, n):
    """ Given an array of k-length positions maps each position
    to n new k+1-length positions by appending numbers from 0 to n-1 
    Example:
    ((),) -> ((0), (2) ... (n-1))
    ((a),(b)) -> ((a, 0), (a,2) ... (a,n-1), (b,0) ... (b,n-1))
    """
    for position in positions:
        for i in range(0, n):
            yield position + (i,);

def menaces(pair1, pair2):
    colK, rowK = pair1;
    col, row = pair2;
    return abs(rowK - row) == abs(colK - col) or rowK == row;

def nQueens(boardSize):
    """ Given a dimension boardSize, returns a filter with all the solutions to
    n-queens problem where n is boardSize"""
    def recur(i):
        if (i == 0): return ((),);
        safeK = lambda x: safe(x, i-1);
        return filter(safeK, addNewColumn(recur(i-1), boardSize));
    return recur(boardSize);

Answer 1

Looks like some pretty spicy (nicely formatted) code :)

I'll explain my steps to make it look better and show you the end result. You can follow along and change them while they pop up.

1. Remove the semicolons :P This is Python, not Javascript or C++ or (insert another language). Though they are optional and are perfectly fine to put in, they can cause clutter and reduce readability.

2. All functions and variables should be lower_case_with_underscores, classes as UpperCase, and constants as UPPER_CASE_WITH_UNDERSCORES. We'll replace them with the correct convention.

3. Replace map and filter with generator expressions. They may be a bit slower, but it'll be more readable as it resembles normal Python syntax such as the for loop and the if statement.

4. That recursive function at the bottom can be changed so that it doesn't need that lambda. It can make it look nicer without the hidden extra arguments.

5. The top function looks pretty complicated... Here's how to change it so that it doesn't need any variables. The first argument to menaces will always be the same, so it can be typed out manually. It's the kth item of enumerate(positions), so the argument is (k, positions[k]). Note: we can remove the import functools as fn line.

6. The second function is a generator function, meaning it returns an iterator. This one is just combining results from two for loops. Nested for loops are discouraged, so let's replace it with itertools.product(positions, range(n)). We'll have to import itertools too. Note: range(0, n) is the same as range(n).

7. I've changed the expression of position + (i,) to (*position, i). I'm not sure which looks better, but personally the second one resembles the tuple that it yields better than the first.

8. The third function could do with a little makeover. The new variables aren't really needed. I've replaced them with col_diff and row_diff which store abs(pair1[n] - pair2[n]) where n is 1 for the rows and 0 for the columns. I've also made it explicitly return True or False and it looks better :)

9. The doc strings also have a convention (for some reason). Single line doc strings can have the triple quotes (""") at the start and end of the line, but multiline doc strings have them on their own lines. Note: I edited the example in add_new_column to look like an interactive session that shows how the function works.

Here is the final result:

import itertools


def safe(positions, k):
    """Given a list of positions, return True if the queen on the `k` column is safe."""
    return not any(
        menaces((k, positions[k]), pair)
        for pair in enumerate(positions)
        if pair[0] != k
    )


def add_new_column(positions, n):
    """
    Given an array of k-length `positions`, map each position to `n` new
    k+1-length positions by appending numbers from 0 to `n` - 1.

    >>> list(add_new_column(((),), 3))
    [(0,), (1,), (2,)]
    >>> list(add_new_column(((1,), (2,)), 3))
    [(1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]
    """
    for position, i in itertools.product(positions, range(n)):
        yield (*position, i)


def menaces(pair1, pair2):
    row_diff = abs(pair1[1] - pair2[1])
    col_diff = abs(pair1[0] - pair2[0])
    if row_diff == col_diff:
        return True
    elif pair1[1] == pair2[1]:
        return True
    else:
        return False


def n_queens(board_size):
    """
    Given a dimension `board_size`, return an iterator with all the solutions
    to n-queens problem where n is `board_size`.
    """
    def _recur(i):
        if not i:
            return ((),)
        else:
            return (
                row
                for row in add_new_column(_recur(i - 1), board_size)
                if safe(row, i - 1)
            )
    return _recur(board_size)

An extra change I would make is to the add_new_column function. It receives the positions first which mean that it has to be in a tuple. By moving the n parameter to the front and using variable length parameters (*positions), the call can look less repetitive.

Here's the new function signature: (You can change the other calls if you like)

def add_new_column(n, *positions):
    """
    Given an array of k-length `positions`, map each position to `n` new
    k+1-length positions by appending numbers from 0 to `n` - 1.

    >>> list(add_new_column(3, ()))
    [(0,), (1,), (2,)]
    >>> list(add_new_column(3, (1,), (2,)))
    [(1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]
    """
    # Here's the code