pythonでNクイーン問題

先週は秋季の応用情報処理技術者試験がありましたね。

午後試験のアルゴリズム問題としてNクイーン問題がとりあげられていたので、紹介されていた処理をそのままpythonで書いてみました。
出題では配列の添字が1から始まってましたが、そこは0からに変えてあります。

ちゃんとした解答はitecですでに公開されてます。
http://www.itec.co.jp/auto_mark/answer/index.html

# coding=utf-8

class NQueen(object):
    def __init__(self, n):
        self.n = n
        self.pos = ['FREE'] * n
        self.col = ['FREE'] * n
        self.upwd = ['FREE'] * ((2 * n) + 1)
        self.downwd  = ['FREE'] * ((2 * n) + 1)

    def search(self, i):
        for k in range(0, self.n):
            if self.col[k] == 'FREE' and self.upwd[i+k] == 'FREE' and self.downwd[self.n+i-k-1] == 'FREE':
                self.pos[i] = k
                self.col[k] = 'NOT_FREE'
                self.upwd[i+k]  = 'NOT_FREE'
                self.downwd[self.n+i-k-1] = 'NOT_FREE'
                if i == self.n-1:
                    return 'SUCCESS'
                else:
                    if (self.search(i+1) == 'SUCCESS'):
                        return 'SUCCESS'
                    else:
                        #print 'クイーンを取り除く'
                        self.pos[i] = 0
                        self.col[k] = 'FREE'
                        self.upwd[i+k]  = 'FREE'
                        self.downwd[self.n+i-k-1] = 'FREE'

        return 'FAILURE'

    def disp(self):
        for i in range(0, self.n):
            for k in range(0, self.n):
                if self.pos[i] == k:
                    print 'Q',
                else:
                    print '-',
            print ""


queen = NQueen(4)

if queen.search(0) == 'SUCCESS':
    queen.disp()
else:
    print '解が見つからなかった'

4*4の出力結果は以下になり、出題の通りになります。

- Q - - 
- - - Q 
Q - - - 
- - Q -

コメントアウトしてありますが、「クイーンを取り除く」の部分の処理が何回実行されるか答えよ、という問題もありました。
コメントを外して、実行してみると4回実行されることがわかります。