N-Queen Solver (某S講義用につくったもの)

nqueen.py

import sys
from collections import deque
INF = 10**9+7
argc = len(sys.argv)
N = int(sys.argv[1]) if argc==2 else int(raw_input())
def tcheck(q, c, idx):
    # (q[i], i) <=> (idx, c)
    return all(q[i]+i!=idx+c and q[i]-i!=idx-c for i in xrange(c))
q_ = [INF]*N
def rotQ(q):
    for i in xrange(N):
        q_[i] = q[i]
        q[i] = INF
    for i in xrange(N):
        # (q[i], i) => (N-1-i, q[i])
        if q_[i]!=INF:
            q[q_[i]] = N-1-i
def mirQ(q):
    for i in xrange(N):
        q[i] = N-1-q[i] if q[i]!=INF else INF
def getMin(q):
    qq = q[:]
    mm = q[:]
    for i in xrange(2):
        for j in xrange(4):
            rotQ(qq)
            if qq<mm:
                for i in xrange(N):
                    mm[i] = qq[i]
        mirQ(qq)
    return mm
def rotateMin(x, y):
    if y==0 or y==N-1:
        return min(x, N-1-x)
    elif x==0 or x==N-1:
        return min(y, N-1-y)
    return N
def getMinPart(q, c):
    qq = q[:c]+[INF]*(N-c)
    mm = q[:]
    for i in xrange(2):
        for j in xrange(4):
            rotQ(qq)
            if qq<mm:
                for i in xrange(N):
                    mm[i] = qq[i]
        mirQ(qq)
    return mm
dup = 0
fall = 0
def dfs(q, c, que):
    global dup, fall
    if c==N:
        qq = getMin(q)
        if qq<q:
            sys.stdout.write("\033[31m")
        for i in xrange(N):
            for j in xrange(N):
                print ("Q" if j==q[i] else "."),
            print  # q[i]
        if qq<q:
            dup += 1
            sys.stdout.write("\033[m")
            print q, qq
        print
        return 1
    ret = 0
    for i in xrange(N-c):
        fall += 1
        t = que.popleft()
        #if (c==0 or tcheck(q, c, t)) and (c==0 or rotateMin(t, c)>=q[0]):
        if(c==0 or tcheck(q, c, t)):
            q[c] = t
            if(q <= getMinPart(q, c+1)):
                ret += dfs(q, c+1, que)
        que.append(t)
    return ret
ans = dfs([0]*N, 0, deque(range(N)))
print ans, "\033[31m"+str(ans-dup)+"\033[m"
print "Count : ", fall