Python N皇后问题实现代码

#首先生成一个棋盘
import numpy as np
from ordered_set import OrderedSet
ndim=4
#检查方法
def check(grid):
    sum=0
    #行循环
    for i in range(grid.shape[0]):
        if(np.sum(grid[i])>1):
            #print("行检查-失败：".format(grid))
            return False
    #列循环
    for i in range(grid.shape[1]):
        if(np.sum(grid[:,i])>1):
            #print("列检查-失败：".format(grid))
            return False
    #循环对角线（从左往右\）
    for i in range(-(ndim-1),ndim):
        eye= np.eye(ndim,ndim,i)
        flip_eye=np.flip(eye,axis=1)
        eye_xy=np.argwhere(eye==1)
        flip_eye_xy=np.argwhere(flip_eye)
        sum=0
        for j in eye_xy:
            sum+=grid[j[0]][j[1]]
            if(sum>1):
                return False
        sum=0
        for j in flip_eye_xy:
            sum+=grid[j[0]][j[1]]
            if(sum>1):
                return False

    return True 
#盲目搜索-枚举
def search():
    grid=np.zeros((ndim,ndim))
    print("{}皇后初始棋盘：".format(ndim))
    print(grid)
    print("--------------------")

    y_index=0
    x_index=0
    #全部的可选择范围
    abc =set(range(ndim))
    #以使用的范围值 (以使用记录需要有序集合)
    a= OrderedSet(set())
    #以尝试的范围值
    b= set()
    count=0
    for i in range(ndim**ndim):
        z=list((abc^a)&(abc^b))
        #print("z:{}--{}".format(z,len(z)))
        if len(z)==0:
            if y_index==0:
                    print("搜索结束，共搜索{}次，共有{}个结果！！！".format(i+1,count))
                    break
            else:
                #回到上一列 (将上一列的使用信息添加到新的尝试集合中，并且删除使用记录)
                y_index-=1
                o=list(a)[y_index]
                b.clear()
                for i in abc:
                    if i<=o:
                        b.add(i)
                a.remove(o)
                grid[:,y_index][o]=0
                continue
        else:
            x_index=z[0]
        
        grid[:,y_index][x_index]=1
        #print(grid)
        #print('----------------')
        if check(grid):
            
            if y_index==ndim-1:
                print("最终结果如下,共尝试{}次,这是第{}个结果".format(i+1,count+1))
                print(grid)
                count+=1
                b.add(x_index)
                grid[:,y_index][x_index]=0
                continue
            else:
                y_index+=1
                b.clear()
            a.append(x_index)
        else:
            #回潮
            #print("回潮前：{}".format(i))
            #print("{}".format(grid))
            b.add(x_index)
            grid[:,y_index][x_index]=0
            if len(list((abc^a)&(abc^b)))==0:
                if y_index==0:
                    print("搜索结束，共搜索{}次，共有{}个结果！！！".format(i+1,count))
                    break
                else:
                    #回到上一列 (将上一列的使用信息添加到新的尝试集合中，并且删除使用记录)
                    y_index-=1
                    o=list(a)[y_index]
                    b.clear()
                    for i in abc:
                        if i<=o:
                            b.add(i)
                    a.remove(o)
                    grid[:,y_index][o]=0
            #print("回潮后：")
            #print("{}".format(grid))    
def test():
    grid=np.zeros((4,4))
    grid[0,2]=1
    grid[1,0]=1
    grid[2,3]=1
    grid[3,1]=1
    print(grid)

    if(check(grid)):
        print("通过测试")
    else :
        print("测试失败")

if __name__ == '__main__':
    #test()
    ndim = int(input('请输入皇后的数量：'))
    search()