{"id":"dsa-backtracking-n-queens","slug":"n-queens","title":"The N queens","text":"A chess queen can move any number of squares along its row, its column or either of its two diagonals, and it attacks every square it could move to. The **N queens** puzzle asks you to place n queens on an n × n board so that none of them attacks another. On a real 8 × 8 chessboard there are 92 ways to do it, and the puzzle is a favorite of interviewers who want to see backtracking.\n\nMaking whole boards and then testing them wastes almost all its time. Eight queens can stand on 64 squares in about 4.4 billion ways. No two queens can share a row, so give each row exactly one: that still leaves 8⁸ = 16,777,216 boards, a column for each of the 8 queens. Ask for different columns too, and you're down to the 8! = 40,320 orderings that `itertools.permutations(range(8))` makes, but n! climbs quickly as well: over 3.6 million boards for n = 10 and almost 40 million for n = 11. Backtracking does better. It fills the board from the top row down, and a queen goes on a square only if no queen placed so far attacks it. A row where every square is attacked means the rows above it can't be part of a solution, so the search goes back up a row and moves that queen along. For n = 8 the whole search takes 2,057 calls.\n\n### Diagonals in O(1)\n\nTo check a square quickly, keep three sets: the columns that have a queen, and the two kinds of diagonals. Along a diagonal that goes down to the right, `row - col` stays the same. Along one that goes down to the left, `row + col` does:\n\n```python\nfor row in range(4):\n    diag_keys = \"\".join(f\"{row - col:3}\" for col in range(4))\n    anti_keys = \"\".join(f\"{row + col:3}\" for col in range(4))\n    print(diag_keys, \"   \", anti_keys)\n```\n\n```output\n  0 -1 -2 -3       0  1  2  3\n  1  0 -1 -2       1  2  3  4\n  2  1  0 -1       2  3  4  5\n  3  2  1  0       3  4  5  6\n```\n\nSo the square `(row, col)` is attacked when `col in cols`, `row - col in diag` or `row + col in anti`: three O(1) lookups in sets, instead of a loop over every queen placed so far.\n\n### The search\n\n```python\ndef place(row, n, queens, cols, diag, anti):\n    if row == n:                       # a queen in every row: a solution\n        for col in queens:\n            print(\".\" * col + \"Q\" + \".\" * (n - col - 1))\n        print()\n        return 1\n    found = 0\n    for col in range(n):\n        if col in cols or row - col in diag or row + col in anti:\n            continue                   # attacked: prune this branch\n        queens.append(col)             # choose\n        cols.add(col)\n        diag.add(row - col)\n        anti.add(row + col)\n        found += place(row + 1, n, queens, cols, diag, anti)   # explore\n        queens.pop()                   # un-choose\n        cols.remove(col)\n        diag.remove(row - col)\n        anti.remove(row + col)\n    return found\n\n\nprint(place(0, 4, [], set(), set(), set()), \"solutions\")\n```\n\n```output\n.Q..\n...Q\nQ...\n..Q.\n\n..Q.\nQ...\n...Q\n.Q..\n\n2 solutions\n```\n\n### How it works\n\n- When `place` is called for a row, every row above it already holds a queen that nothing attacks. So `row == n` means the board is full: one more solution.\n- Follow n = 4 from a queen in the top-left corner. The second row's queen can go in column 2 or 3. After column 2, every square of the third row is attacked. After column 3, the third row's queen fits in column 1, and then every square of the fourth row is attacked. Both ways are dead, so the search takes the first queen back and moves it to column 1, where it finds the first solution. Going back up like this, to change a choice made earlier, is what the word **backtracking** describes.\n- The choose changes four things, the list and the three sets, and the un-choose must undo every one of them. Forget one, say `diag.remove`, and that diagonal stays blocked for the rest of the search: the count comes out too small, with no error to warn you.\n\n### What it costs\n\nSolutions multiply as the board grows: 352 for n = 9, 724 for n = 10, 2,680 for n = 11 and 14,200 for n = 12. The search is still exponential, but pruning cuts the tree down to a sliver: for n = 10 it makes 35,539 calls instead of checking 10¹⁰ boards. Each call does O(n) work, one O(1) check per column, and the recursion is only n + 1 calls deep, so the extra memory is O(n): the sets and the list of queens. Checking a new queen against every queen placed so far, with a loop, also works; it's O(n) per check instead of O(1), a few times slower. What makes the search fast is checking **before** placing, so that dead boards are never finished.\n","task":"A 4 × 4 board has queens at (0, 1) and (1, 3). What does this check of the third row print? Type each line it prints.\n","kind":"predict","seed":"cols = {1, 3}\ndiag = {0 - 1, 1 - 3}\nanti = {0 + 1, 1 + 3}\nrow = 2\nfor col in range(4):\n    if col in cols or row - col in diag or row + col in anti:\n        print(col, \"attacked\")\n    else:\n        print(col, \"free\")\n","solution":"cols = {1, 3}\ndiag = {0 - 1, 1 - 3}\nanti = {0 + 1, 1 + 3}\nrow = 2\nfor col in range(4):\n    if col in cols or row - col in diag or row + col in anti:\n        print(col, \"attacked\")\n    else:\n        print(col, \"free\")\n","hints":["`diag` holds -1 and -2, and `anti` holds 1 and 4. For each column, work out `row - col` and `row + col` with `row = 2`.","Columns 1 and 3 already have a queen. Column 2 gives `row + col` = 4, which is in `anti`: the queen at (1, 3) attacks it along a diagonal.","Column 0 gives 2 for both sums, and 2 is in neither set."],"tests":[{"name":"Output","stdin":"","expect":"0 free\n1 attacked\n2 attacked\n3 attacked","hidden":false}],"require":[],"forbid":[],"lines":["0 free","1 attacked","2 attacked","3 attacked"],"more":[{"task":"Write `count_queens(n)`, which returns the number of ways to put n queens on an n × n board so that no two attack each other. Don't print the boards. It must be fast for n = 11, so build the board row by row and skip attacked squares with the three sets.\n\nFor example, `print(count_queens(4), count_queens(6))` prints:\n\n```\n2 4\n```\n","kind":"code","seed":"def count_queens(n):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(count_queens(4), count_queens(6))\n","solution":"def place(row, n, cols, diag, anti):\n    if row == n:\n        return 1\n    found = 0\n    for col in range(n):\n        if col in cols or row - col in diag or row + col in anti:\n            continue\n        cols.add(col)\n        diag.add(row - col)\n        anti.add(row + col)\n        found += place(row + 1, n, cols, diag, anti)\n        cols.remove(col)\n        diag.remove(row - col)\n        anti.remove(row + col)\n    return found\n\n\ndef count_queens(n):\n    return place(0, n, set(), set(), set())\n\n\nif __name__ == \"__main__\":\n    print(count_queens(4), count_queens(6))\n","hints":["Write a helper `place(row, n, cols, diag, anti)` that returns the number of solutions for the rows from `row` on. When `row == n`, every row has a queen: return 1.","Loop over the columns. Skip a column when `col in cols or row - col in diag or row + col in anti`. Otherwise add the three keys, add the count from `row + 1`, and remove the three keys again.","`count_queens(n)` calls the helper with row 0 and three empty sets. Filling every row first and checking the board only at the end gives the right count too, but for n = 11 that's 11¹¹ boards, about 285 billion."],"tests":[{"name":"print(count_queens(4), count_queens(6))","stdin":"","call":"print(count_queens(4), count_queens(6))\n","expect":"2 4","hidden":false},{"name":"count_queens(1)","stdin":"","call":"count_queens(1)\n","expect":"1","hidden":true,"returns":true},{"name":"count_queens(2)","stdin":"","call":"count_queens(2)\n","expect":"0","hidden":true,"returns":true},{"name":"count_queens(3)","stdin":"","call":"count_queens(3)\n","expect":"0","hidden":true,"returns":true},{"name":"count_queens(5)","stdin":"","call":"count_queens(5)\n","expect":"10","hidden":true,"returns":true},{"name":"count_queens(11)","stdin":"","call":"count_queens(11)\n","expect":"2680","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[]},{"task":"Often one solution is enough. Write `first_solution(n)`: the first solution the search finds when it tries the columns from left to right, as a list of the queens' columns, row by row, or `None` when there's no solution. Stop as soon as you have one: make the search return `True` when the last row is placed, and return at once all the way up, without undoing anything. It must be fast for n = 21, which has hundreds of billions of solutions.\n\nFor example, `first_solution(4)` returns:\n\n```\n[1, 3, 0, 2]\n```\n","kind":"code","seed":"def first_solution(n):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(first_solution(4))\n","solution":"def search(row, n, queens, cols, diag, anti):\n    if row == n:\n        return True\n    for col in range(n):\n        if col in cols or row - col in diag or row + col in anti:\n            continue\n        queens.append(col)\n        cols.add(col)\n        diag.add(row - col)\n        anti.add(row + col)\n        if search(row + 1, n, queens, cols, diag, anti):\n            return True           # found: keep the queens as they are\n        queens.pop()\n        cols.remove(col)\n        diag.remove(row - col)\n        anti.remove(row + col)\n    return False\n\n\ndef first_solution(n):\n    queens = []\n    if search(0, n, queens, set(), set(), set()):\n        return queens\n    return None\n\n\nif __name__ == \"__main__\":\n    print(first_solution(4))\n","hints":["Start from the lesson's `place`, with a list `queens` of the columns. Make the helper return `True` when `row == n`, and `False` after its loop has tried every column.","After placing a queen, write `if search(row + 1, ...): return True`. When the rest of the board worked, return straight away: the queens stay in the list, which is the answer.","Only when the call below returns `False` do you undo the queen and try the next column. A search that goes on after the first solution and keeps it gives the right answer too, but for n = 21 it never ends."],"tests":[{"name":"first_solution(4)","stdin":"","call":"first_solution(4)\n","expect":"[1, 3, 0, 2]","hidden":false,"returns":true},{"name":"first_solution(1)","stdin":"","call":"first_solution(1)\n","expect":"[0]","hidden":true,"returns":true},{"name":"print(first_solution(3))","stdin":"","call":"print(first_solution(3))\n","expect":"None","hidden":true},{"name":"first_solution(6)","stdin":"","call":"first_solution(6)\n","expect":"[1, 3, 5, 0, 2, 4]","hidden":true,"returns":true},{"name":"first_solution(8)","stdin":"","call":"first_solution(8)\n","expect":"[0, 4, 7, 5, 2, 6, 1, 3]","hidden":true,"returns":true},{"name":"first_solution(16)","stdin":"","call":"first_solution(16)\n","expect":"[0, 2, 4, 1, 12, 8, 13, 11, 14, 5, 15, 6, 3, 10, 7, 9]","hidden":true,"big":true,"returns":true},{"name":"first_solution(21)","stdin":"","call":"first_solution(21)\n","expect":"[0, 2, 4, 1, 3, 8, 10, 14, 20, 17, 19, 16, 18, 6, 11, 9, 7, 5, 13, 15, 12]","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[]}]}