hw(AI/2): Remove max_queens logic, add tests for methods in cell

3 files changed, 167 insertions, 57 deletions

diff --git a/Year 4/Artificial Intelligence/hw02/src/cell.rs b/Year 4/Artificial Intelligence/hw02/src/cell.rs
index 9b36199..8df0589 100644
--- a/Year 4/Artificial Intelligence/hw02/src/cell.rs
+++ b/Year 4/Artificial Intelligence/hw02/src/cell.rs
@@ -2,17 +2,28 @@ use super::Puzzle;
 use rand::Rng;
 
 impl Puzzle {
+    /// O(N)
     pub fn max_queen_col(&mut self) -> usize {
-        if self.max_attack == 3 {
-            return self.N;
+        let (mut max_attack, mut max_queens) = (i32::MIN, Vec::new());
+        for col in 0..self.N {
+            let attacks = self.attacks(col, self.in_col[col]);
+
+            if attacks > max_attack {
+                max_attack = attacks;
+                max_queens.clear();
+            }
+            if attacks == max_attack {
+                max_queens.push(col);
+            }
         }
 
-        if self.max_queens.len() > 1 {
-            return self.max_queens[self.rng.gen_range(0..self.max_queens.len())];
+        if max_queens.len() > 1 {
+            return max_queens[self.rng.gen_range(0..max_queens.len())];
         }
-        return self.max_queens[0];
+        return max_queens[0];
     }
 
+    /// O(N)
     pub fn min_col_row(&mut self, colI: usize) -> usize {
         let (mut min_row, mut min_ind) = (i32::MAX, Vec::new());
         for row in 0..self.N {
@@ -33,37 +44,21 @@ impl Puzzle {
         return min_ind[0];
     }
 
+    /// O(1)
     pub fn update(&mut self, colI: usize, val: i32) {
-        let old_attack = self.attacks(colI, self.in_col[colI]);
+        assert!(colI < self.N && (val == 1 || val == -1), "Bad arguments to update: {} {}!", colI, val);
 
         self.per_row[self.in_col[colI]]                    += val;
         self.per_d1[colI + self.N - 1 - self.in_col[colI]] += val;
         self.per_d2[colI + self.in_col[colI]]              += val;
 
         if val < 0 {
-            if old_attack != self.max_attack { return; }
-
-            for i in 0..self.max_queens.len() {
-                if self.max_queens[i] == colI {
-                    self.max_queens.remove(i);
-                    return;
-                }
-            }
-            unreachable!();
-        }
-
-        let current_attack = self.attacks(colI, self.in_col[colI]);
-
-        if self.max_attack < current_attack {
-            self.max_queens.clear();
-            self.max_queens.push(colI);
-            self.max_attack = current_attack;
-        }
-        if self.max_attack == current_attack {
-            self.max_queens.push(colI);
+            // Sanity check when removing a queen
+            self.in_col[colI] = self.N;
         }
     }
 
+    /// O(1)
     pub fn move_queen(&mut self, colI: usize, rowI: usize) {
         if rowI == self.in_col[colI] { return; }
 
@@ -72,3 +67,69 @@ impl Puzzle {
         self.update(colI, 1);
     }
 }
+
+#[test]
+fn test_max_queen_col() {
+    let mut e1 = Puzzle::example1();
+    assert_eq!(e1.max_queen_col(), 2, "Bad max_queen_col on single max!");
+
+    let mut e2 = Puzzle::example2();
+    let mut got = [false; 4];
+    for _ in 0..10000 {
+        got[e2.max_queen_col()] = true;
+    }
+    assert!(got.iter().all(|&x| x), "Bad max_queen_col on multiple max!");
+}
+
+#[test]
+fn test_min_col_row() {
+    let mut e1 = Puzzle::example1();
+    let mut got = [false; 4];
+    for _ in 0..10000 {
+        got[e1.min_col_row(0)] = true;
+    }
+    assert!(got[0] && !got[1] && got[2] && got[3], "Bad min_col_row on multiple (3) min!");
+    got = [false; 4];
+    for _ in 0..10000 {
+        got[e1.min_col_row(2)] = true;
+    }
+    assert!(got[0] && !got[1] && got[2] && !got[3], "Bad min_col_row on multiple (2) min!");
+
+    let mut e2 = Puzzle::example2();
+    assert_eq!(e2.min_col_row(0), 3, "Bad min_col_row on single min!");
+    assert_eq!(e2.min_col_row(3), 1, "Bad min_col_row on single min!");
+}
+
+#[test]
+fn test_update() {
+    let mut e1 = Puzzle::example1();
+    e1.update(0, -1);
+    e1.update(2, -1);
+    e1.in_col[0] = 3;
+    e1.update(0, 1);
+    e1.in_col[2] = 1;
+    e1.update(2, 1);
+    let mut e1s = Puzzle::example1();
+    assert_eq!(e1, e1s, "Bad update!");
+}
+
+#[test]
+fn test_move_queen() {
+    let mut e1 = Puzzle::example1();
+    e1.move_queen(0, 0);
+    e1.move_queen(0, 3);
+    let mut e1s = Puzzle::example1();
+    assert_eq!(e1, e1s, "Bad move_queen identity!");
+
+    let mut e1 = Puzzle::example1();
+    e1.move_queen(2, 0);
+    e1.move_queen(2, 1);
+    let mut e1s = Puzzle::example1();
+    assert_eq!(e1, e1s, "Bad move_queen identity on moved max!");
+
+    let mut e1 = Puzzle::example1();
+    e1.move_queen(0, 0);
+    e1.move_queen(2, 2);
+    let mut e2 = Puzzle::example2();
+    assert_eq!(e1, e2, "Bad move_queen on changed maxes!");
+}
diff --git a/Year 4/Artificial Intelligence/hw02/src/helper.rs b/Year 4/Artificial Intelligence/hw02/src/helper.rs
index 478e400..164184f 100644
--- a/Year 4/Artificial Intelligence/hw02/src/helper.rs
+++ b/Year 4/Artificial Intelligence/hw02/src/helper.rs
@@ -8,9 +8,6 @@ impl Puzzle {
             per_d1:  [0; 2 * MAXN - 1],
             per_d2:  [0; 2 * MAXN - 1],
 
-            max_attack: 0,
-            max_queens: Vec::with_capacity(N),
-
             N: N,
             rng: rand::thread_rng(),
         }
@@ -23,18 +20,17 @@ impl Puzzle {
                 self.per_row[i] = 0;
             }
             else if i < MAXN {
-                assert_eq!(self.in_col[i], 0);
-                assert_eq!(self.per_row[i], 0);
+                assert_eq!(self.in_col[i], 0, "Someone wrote to in_col at index {}!", i);
+                assert_eq!(self.per_row[i], 0, "Someone wrote to per_row at index {}!", i);
             }
             self.per_d1[i] = 0;
             self.per_d2[i] = 0;
         }
-        self.max_attack = 0;
-        self.max_queens.clear();
     }
 
+    /// O(1)
     pub fn attacks(&self, col: usize, row: usize) -> i32 {
-        assert!(col < self.N && row < self.N);
+        assert!(col < self.N && row < self.N, "Bad attacks index: {} {}!", col, row);
 
         1 +                 // Always a queen on a column
         self.per_row[row] + // Queens on the row
@@ -44,10 +40,10 @@ impl Puzzle {
         // If we're in the cell with a queen, remove repeated counts of it
     }
 
-    // | |*| | |
-    // | | |*| |
-    // | | | |*|
-    // |*| | | |
+    // | |*| | |    |2|3|2|4|
+    // | | |*| |    |3|2|4|2|
+    // | | | |*|    |2|4|2|3|
+    // |*| | | |    |2|2|3|2|
     pub fn example1() -> Puzzle {
         let mut r = Puzzle::new(4);
         r.in_col[0] = 3;
@@ -62,6 +58,30 @@ impl Puzzle {
         r.per_d2[1] = 1;
         r.per_d2[3] = 2;
         r.per_d2[5] = 1;
+
+        r
+    }
+
+    // |*|*| | |    |3|3|3|3|
+    // | | | | |    |2|3|3|2|
+    // | | |*|*|    |3|3|3|3|
+    // | | | | |    |1|2|2|3|
+    pub fn example2() -> Puzzle {
+        let mut r = Puzzle::new(4);
+        r.in_col[0] = 0;
+        r.in_col[1] = 0;
+        r.in_col[2] = 2;
+        r.in_col[3] = 2;
+        r.per_row[0] = 2; r.per_row[1] = 0; r.per_row[2] = 2; r.per_row[3] = 0;
+
+        r.per_d1[3] = 2;
+        r.per_d1[4] = 2;
+
+        r.per_d2[0] = 1;
+        r.per_d2[1] = 1;
+        r.per_d2[4] = 1;
+        r.per_d2[5] = 1;
+
         r
     }
 }
@@ -69,23 +89,44 @@ impl Puzzle {
 #[test]
 fn test_attacks() {
     let e1 = Puzzle::example1();
-    assert_eq!(e1.attacks(0, 0), 2); // top-left
-    assert_eq!(e1.attacks(0, 1), 3); // below top-left
-    assert_eq!(e1.attacks(0, 2), 2);
-    assert_eq!(e1.attacks(0, 3), 2);
+    assert_eq!(e1.attacks(0, 0), 2, "Bad attacks!"); // top-left
+    assert_eq!(e1.attacks(0, 1), 3, "Bad attacks!"); // below top-left
+    assert_eq!(e1.attacks(0, 2), 2, "Bad attacks!");
+    assert_eq!(e1.attacks(0, 3), 2, "Bad attacks!");
+
+    assert_eq!(e1.attacks(1, 0), 3, "Bad attacks!"); // right of top-left
+    assert_eq!(e1.attacks(1, 1), 2, "Bad attacks!");
+    assert_eq!(e1.attacks(1, 2), 4, "Bad attacks!");
+    assert_eq!(e1.attacks(1, 3), 2, "Bad attacks!");
+
+    assert_eq!(e1.attacks(2, 0), 2, "Bad attacks!");
+    assert_eq!(e1.attacks(2, 1), 4, "Bad attacks!");
+    assert_eq!(e1.attacks(2, 2), 2, "Bad attacks!");
+    assert_eq!(e1.attacks(2, 3), 3, "Bad attacks!");
+
+    assert_eq!(e1.attacks(3, 0), 4, "Bad attacks!");
+    assert_eq!(e1.attacks(3, 1), 2, "Bad attacks!");
+    assert_eq!(e1.attacks(3, 2), 3, "Bad attacks!");
+    assert_eq!(e1.attacks(3, 3), 2, "Bad attacks!");
+
+    let e2 = Puzzle::example2();
+    assert_eq!(e2.attacks(0, 0), 3, "Bad attacks!"); // top-left
+    assert_eq!(e2.attacks(0, 1), 2, "Bad attacks!"); // below top-left
+    assert_eq!(e2.attacks(0, 2), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(0, 3), 1, "Bad attacks!");
 
-    assert_eq!(e1.attacks(1, 0), 3); // right of top-left
-    assert_eq!(e1.attacks(1, 1), 2);
-    assert_eq!(e1.attacks(1, 2), 4);
-    assert_eq!(e1.attacks(1, 3), 2);
+    assert_eq!(e2.attacks(1, 0), 3, "Bad attacks!"); // right of top-left
+    assert_eq!(e2.attacks(1, 1), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(1, 2), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(1, 3), 2, "Bad attacks!");
 
-    assert_eq!(e1.attacks(2, 0), 2);
-    assert_eq!(e1.attacks(2, 1), 4);
-    assert_eq!(e1.attacks(2, 2), 2);
-    assert_eq!(e1.attacks(2, 3), 3);
+    assert_eq!(e2.attacks(2, 0), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(2, 1), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(2, 2), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(2, 3), 2, "Bad attacks!");
 
-    assert_eq!(e1.attacks(3, 0), 4);
-    assert_eq!(e1.attacks(3, 1), 2);
-    assert_eq!(e1.attacks(3, 2), 3);
-    assert_eq!(e1.attacks(3, 3), 2);
+    assert_eq!(e2.attacks(3, 0), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(3, 1), 2, "Bad attacks!");
+    assert_eq!(e2.attacks(3, 2), 3, "Bad attacks!");
+    assert_eq!(e2.attacks(3, 3), 3, "Bad attacks!");
 }
diff --git a/Year 4/Artificial Intelligence/hw02/src/main.rs b/Year 4/Artificial Intelligence/hw02/src/main.rs
index 2de0630..970a04c 100644
--- a/Year 4/Artificial Intelligence/hw02/src/main.rs
+++ b/Year 4/Artificial Intelligence/hw02/src/main.rs
@@ -6,19 +6,27 @@ use std::time::Instant;
 
 pub const MAXN: usize = 10010;
 
+#[derive(Debug)]
 pub struct Puzzle {
     in_col:  [usize; MAXN],       // Позиция в колона
     per_row: [i32; MAXN],         // Бройка на ред
     per_d1:  [i32; 2 * MAXN - 1], // Бройка на главен диагонал
     per_d2:  [i32; 2 * MAXN - 1], // Бройка на втори диагонал
 
-    max_attack: i32,
-    max_queens: Vec<usize>,
-
     N: usize,
     rng: ThreadRng,
 }
 
+impl PartialEq for Puzzle {
+    fn eq(&self, other: &Puzzle) -> bool {
+        (self.in_col == other.in_col) &&
+        (self.per_row == other.per_row) &&
+        (self.per_d1 == other.per_d1) &&
+        (self.per_d2 == other.per_d2) &&
+        (self.N == other.N)
+    }
+}
+
 mod helper;
 mod cell;
 mod alg;