N-Queen Problem: CPython 2.6.5 vs PyPy 1.5.0

It is amazing how small and compress a Python code can be for solving 8-Queen problem. Using Python’s itertools to compute permutations and yield keyword for generators, all N-Queen problem comes to mere 8 lines of code:

def hit(p):
    for i in range(len(p) - 1):
        for j in range(i + 1, len(p)):
            if j - i == abs(p[i] - p[j]): return True
    return False

def n_queen(n):
    for p in itertools.permutations(range(n)):
        if not hit(p): yield p

One may add just a bit more code to print the board like this:

base_char = ord('A')

def print_board(s):
    end = len(s)
    board = ['   '.join([str(i) for i in range(1, end + 1)]) + '\n']
    for i, r in enumerate(s):
        board += "%s: %s * %s\n" % (chr(base_char+i), (' - ' * r), (' - ' * (end - r - 1)))
    return ''.join(board)

def solve(n):
    count = 0
    for s in n_queen(n):
        count += 1
        print "==== solution %d =====" % count
        print s
        print print_board(s)
    return count

I used this code snippet to compare Python 2.6.5 with PyPy 1.5.0 JIT. It is interesting that for this CPU intensive problem, Python does not use all 100% CPU in User space and only about 30% of CPU usage is spent in User space and the rest 70% for System and Nice while PyPy dedicated most of CPU cycles to User CPU and hence finished up much faster.

Here are my results in an old single core 1.7G Pentium M processor. I expect PyPy JIT to do even better in a modern SSE-2 multi-core processor.

# board count cpython2.6.5 pypy1.5
5 10 0.012 0.007
6 4 0.048 0.064
7 40 0.327 0.224
8 92 1.556 0.562
9 352 15.207 3.685
10 724 156.736 37.676
11 2680 1723.404 427.216

Bellow is the chart comparing the performance upon above benchmark:

Next I tried to speed up the code by making a minor change in hit() method. To prevent len(p) call I passed the length as argument and replaced abs() call with a check for both positive and negative results. These two simple changes as bellow made the code almost 3 times faster:

def hit(p, n):
    for i in range(n - 1):
        for j in range(i + 1, n):
            d = p[i] - p[j]
            if j - i == d or i - j == d: return True
    return False

def n_queen(n):
    for p in itertools.permutations(range(n)):
        if not hit(p, n): yield p

I was expecting the PyPy fail to speed up the execution times as it did for basic code but amazingly the second code’s speed up is even better in PyPy.

Update 31 May 2011 :
Just one update which seemed necessary, back to some comments to this post and also in YCombinator news forum, it is obvious that what I proposed above was not an optimal solution for N-Queen. In fact N-Queen is a classic backtracking problem and can be solved much efficiently using BT but what I tried to show was 1st the power of Python generators and itertools to solve this problem is a few lines (while preserving simplicity) and 2nd using this CPU intensive brute-force search to compare CPython and PyPy implementations.

So I decided to mention a backtracking implementation which solves the problem much much faster and easily covers boards of size up to 14.

def last_one_hits(answers, size):
    i = size - 1
    x, j = answers[i], 0
    for y in answers[:i]:
        d = x - y
        if x == y or i - j == d or j - i == d: return True
        j += 1
    return False

def n_queen_bt(n):
    found = [0] * n # initiate an answer of size n
    size = 1 # in start, we have 1 queen in cloumn/row (0,0)

    while size > 0:
        hits = last_one_hits(found, size) # new queen hits prevs

        if size == n and not hits: #full and does not hit
            yield found

        if size < n and not hits: #not full or hit, add a new
            found[size] = 0
            size += 1
        else:
            while size > 0: #move last or maybe the ones before
                found[size-1] += 1
                if found[size-1] == n: #end of this clmn,go back
                    size -= 1
                else:
                    break #did the movement job, let's continue