Demo/scripts/queens.py

   1 #! /usr/bin/env python
   2
   3 """N queens problem.
   4
   5 The (well-known) problem is due to Niklaus Wirth.
   6
   7 This solution is inspired by Dijkstra (Structured Programming).  It is
   8 a classic recursive backtracking approach.
   9
  10 """
  11
  12 N = 8                                   # Default; command line overrides
  13
  14 class Queens:
  15
  16     def __init__(self, n=N):
  17         self.n = n
  18         self.reset()
  19
  20     def reset(self):
  21         n = self.n
  22         self.y = [None] * n             # Where is the queen in column x
  23         self.row = [0] * n              # Is row[y] safe?
  24         self.up = [0] * (2*n-1)         # Is upward diagonal[x-y] safe?
  25         self.down = [0] * (2*n-1)       # Is downward diagonal[x+y] safe?
  26         self.nfound = 0                 # Instrumentation
  27
  28     def solve(self, x=0):               # Recursive solver
  29         for y in range(self.n):
  30             if self.safe(x, y):
  31                 self.place(x, y)
  32                 if x+1 == self.n:
  33                     self.display()
  34                 else:
  35                     self.solve(x+1)
  36                 self.remove(x, y)
  37
  38     def safe(self, x, y):
  39         return not self.row[y] and not self.up[x-y] and not self.down[x+y]
  40
  41     def place(self, x, y):
  42         self.y[x] = y
  43         self.row[y] = 1
  44         self.up[x-y] = 1
  45         self.down[x+y] = 1
  46
  47     def remove(self, x, y):
  48         self.y[x] = None
  49         self.row[y] = 0
  50         self.up[x-y] = 0
  51         self.down[x+y] = 0
  52
  53     silent = 0                          # If true, count solutions only
  54
  55     def display(self):
  56         self.nfound = self.nfound + 1
  57         if self.silent:
  58             return
  59         print '+-' + '--'*self.n + '+'
  60         for y in range(self.n-1, -1, -1):
  61             print '|',
  62             for x in range(self.n):
  63                 if self.y[x] == y:
  64                     print "Q",
  65                 else:
  66                     print ".",
  67             print '|'
  68         print '+-' + '--'*self.n + '+'
  69
  70 def main():
  71     import sys
  72     silent = 0
  73     n = N
  74     if sys.argv[1:2] == ['-n']:
  75         silent = 1
  76         del sys.argv[1]
  77     if sys.argv[1:]:
  78         n = int(sys.argv[1])
  79     q = Queens(n)
  80     q.silent = silent
  81     q.solve()
  82     print "Found", q.nfound, "solutions."
  83
  84 if __name__ == "__main__":
  85     main()