python ch-2

1 files changed, 27 insertions, 0 deletions

diff --git a/challenge-062/sangeet-kar/python/ch-2.py b/challenge-062/sangeet-kar/python/ch-2.py
new file mode 100755
index 0000000000..e243087412
--- /dev/null
+++ b/challenge-062/sangeet-kar/python/ch-2.py
@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+
+import sys
+
+def n_queens_3d (n = 2):
+    solutions = []
+    place_queen ([(i, j, k) for i in range(n) for j in range(n) for k in range(n)], [], solutions)
+    return indices_to_array(max(solutions, key=len), n);
+
+def place_queen (indices, queens, solutions):
+    if not indices: solutions.append(queens)
+    for pos in indices:
+        place_queen ([index for index in indices if is_available(pos, index)], [*queens, pos], solutions)
+
+def is_available(ref, pos):
+    diff = {abs(i - j) for i, j in zip (ref, pos)}
+    return not ( len(diff) < 2 or (len(diff) == 2 and 0 in diff))
+
+def indices_to_array (indices, n):
+    array = [[[0 for _ in range(n)] for _ in range(n)] for _ in range(n)]
+    for i, j, k in indices:
+        array[i][j][k] = 1
+    return array
+
+n = int(sys.argv[1]) if len(sys.argv) > 1 else 2
+print(n_queens_3d (n))
+