algorithms_edu/problems/backtracking/
nqueens.rs

1//! The n-queens puzzle is the problem of placing n queens on an n x n chessboard such that no two queens attack each other.
2//!
3//! Given an integer n, return all distinct solutions to the n-queens puzzle.
4//!
5//! Each solution contains a distinct board configuration of the n-queens' placement, where 'Q' and '.' both indicate a queen and an empty space, respectively.
6//!
7//! # See also
8//!
9//! - [Leetcode](https://leetcode.com/problems/n-queens/)
10
11use std::collections::HashSet;
12
13pub fn solve_n_queens(n: i32) -> Vec<Vec<String>> {
14    Board::new(n as usize).solve()
15}
16
17struct Board {
18    matrix: Vec<Vec<char>>,
19    n: usize,
20    solutions: HashSet<Vec<String>>,
21}
22
23impl Board {
24    pub fn new(n: usize) -> Self {
25        Self {
26            matrix: vec![vec!['.'; n]; n],
27            n,
28            solutions: HashSet::new(),
29        }
30    }
31    pub fn solve(mut self) -> Vec<Vec<String>> {
32        self._solve(0, 0);
33        self.solutions.into_iter().collect()
34    }
35    fn _solve(&mut self, i: usize, count: usize) {
36        if count == self.n {
37            self.add_solution();
38        } else if i == self.n {
39        } else {
40            for col in 0..self.n {
41                if self.safe(i, col) {
42                    self.matrix[i][col] = 'Q';
43                    self._solve(i + 1, count + 1);
44                    self.matrix[i][col] = '.';
45                }
46            }
47        }
48    }
49    fn add_solution(&mut self) {
50        self.solutions.insert(
51            self.matrix
52                .iter()
53                .map(|x| x.iter().copied().collect())
54                .collect(),
55        );
56    }
57
58    fn safe(&self, i: usize, j: usize) -> bool {
59        for j_ in 0..self.n {
60            if self.matrix[i][j_] == 'Q' {
61                return false;
62            }
63        }
64        for i_ in 0..self.n {
65            if self.matrix[i_][j] == 'Q' {
66                return false;
67            }
68        }
69        let (mut i_, mut j_) = (i + 1, j + 1);
70        while i_ > 0 && j_ > 0 {
71            i_ -= 1;
72            j_ -= 1;
73            if self.matrix[i_][j_] == 'Q' {
74                return false;
75            }
76        }
77        let (mut i_, mut j_) = (i, j + 1);
78        while i_ < self.n && j_ > 0 {
79            j_ -= 1;
80            if self.matrix[i_][j_] == 'Q' {
81                return false;
82            }
83            i_ += 1;
84        }
85        let (mut i_, mut j_) = (i, j);
86        while i_ < self.n && j_ < self.n {
87            if self.matrix[i_][j_] == 'Q' {
88                return false;
89            }
90            i_ += 1;
91            j_ += 1;
92        }
93        let (mut i_, mut j_) = (i + 1, j);
94        while i_ > 0 && j_ < self.n {
95            i_ -= 1;
96            if self.matrix[i_][j_] == 'Q' {
97                return false;
98            }
99            j_ += 1;
100        }
101        true
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    #[test]
109    fn test_n_queens() {
110        let n1 = solve_n_queens(1);
111        assert_eq!(n1, vec![vec!["Q".to_string()]]);
112
113        let n2 = solve_n_queens(2);
114        assert!(n2.is_empty());
115
116        let n3 = solve_n_queens(3);
117        assert!(n3.is_empty());
118
119        let mut n4 = solve_n_queens(4);
120        n4.sort();
121        assert_eq!(
122            n4,
123            make_solution(&[
124                &["..Q.", "Q...", "...Q", ".Q.."],
125                &[".Q..", "...Q", "Q...", "..Q."],
126            ])
127        );
128
129        let mut n5 = solve_n_queens(5);
130        let mut n5_expected = make_solution(&[
131            &["..Q..", "....Q", ".Q...", "...Q.", "Q...."],
132            &["...Q.", "Q....", "..Q..", "....Q", ".Q..."],
133            &["....Q", ".Q...", "...Q.", "Q....", "..Q.."],
134            &["Q....", "...Q.", ".Q...", "....Q", "..Q.."],
135            &[".Q...", "....Q", "..Q..", "Q....", "...Q."],
136            &["....Q", "..Q..", "Q....", "...Q.", ".Q..."],
137            &[".Q...", "...Q.", "Q....", "..Q..", "....Q"],
138            &["..Q..", "Q....", "...Q.", ".Q...", "....Q"],
139            &["...Q.", ".Q...", "....Q", "..Q..", "Q...."],
140            &["Q....", "..Q..", "....Q", ".Q...", "...Q."],
141        ]);
142        n5.sort();
143        n5_expected.sort();
144        assert_eq!(n5, n5_expected);
145
146        let n8 = solve_n_queens(8);
147        assert_eq!(n8.len(), 92);
148    }
149
150    fn make_solution(sol: &[&[&'static str]]) -> Vec<Vec<String>> {
151        sol.iter()
152            .map(|&x| x.iter().map(|&s| s.to_owned()).collect())
153            .collect()
154    }
155}