hw(4/AI): Final changes and improvements

1 files changed, 104 insertions, 58 deletions

diff --git a/Year 4/Artificial Intelligence/hw02/src/main.rs b/Year 4/Artificial Intelligence/hw02/src/main.rs
index 0d2ec11..f36f5c6 100644
--- a/Year 4/Artificial Intelligence/hw02/src/main.rs
+++ b/Year 4/Artificial Intelligence/hw02/src/main.rs
@@ -4,56 +4,74 @@ use rand::{Rng, rngs::ThreadRng};
 use std::fmt::{Display, Formatter, Result};
 use std::time::Instant;
 
+const MAXN: usize = 10000;
+
 struct Puzzle {
-    in_col: Vec<usize>,
-    per_row: Vec<i32>,
-    per_d1: Vec<i32>,
-    per_d2: Vec<i32>,
+    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 Puzzle {
+    /* Helper */
+
     pub fn new(N: usize) -> Puzzle {
         Puzzle {
-            in_col:  vec![0; N],
-            per_row: vec![0; N],
-            per_d1:  vec![0; 2 * N - 1],
-            per_d2:  vec![0; 2 * N - 1],
+            in_col:  [0; MAXN],
+            per_row: [0; MAXN],
+            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(),
         }
     }
 
-    fn max_queen_col(&mut self) -> usize {
-        let (mut max_attacks, mut ind) = (0, Vec::new());
-        for col in 0..self.N {
-            let attacks = self.per_row[self.in_col[col]] +
-                          self.per_d1[col + self.N - 1 - self.in_col[col]] +
-                          self.per_d2[col + self.in_col[col]];
+    fn attacks(&self, col: usize, row: usize) -> i32 {
+        self.per_row[row] +
+        self.per_d1[col + self.N - 1 - row] +
+        self.per_d2[col + row]
+    }
 
-            if max_attacks < attacks {
-                max_attacks = attacks;
-                ind.clear();
-            }
-            if max_attacks == attacks {
-                ind.push(col);
+    fn clear(&mut self) {
+        self.max_attack = 0;
+        self.max_queens.clear();
+        for i in 0..2 * self.N - 1 {
+            self.per_d1[i] = 0;
+            self.per_d2[i] = 0;
+            if i < self.N {
+                self.per_row[i] = 0;
             }
         }
+    }
+
+    /* Cells */
+
+    fn max_queen_col(&mut self) -> usize {
+        if self.max_attack == 3 {
+            return self.N;
+        }
 
-        if ind.len() > 1 {
-            return ind[self.rng.gen_range(0..ind.len())];
+        if self.max_queens.len() > 1 {
+            return self.max_queens[self.rng.gen_range(0..self.max_queens.len())];
         }
-        return ind[0];
+        return self.max_queens[0];
     }
 
     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 {
-            let attacks = self.per_row[row] +
-                          self.per_d1[colI + self.N - 1 - row] +
-                          self.per_d2[colI + row];
+            let attacks = self.attacks(colI, row);
 
             if attacks < min_row {
                 min_row = attacks;
@@ -71,10 +89,36 @@ impl Puzzle {
     }
 
     fn update(&mut self, colI: usize, val: i32) {
-        self.per_row[self.in_col[colI]] += val;
+        let old_attack = self.attacks(colI, self.in_col[colI]);
+
+        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;
+        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);
+        }
     }
+
     fn move_queen(&mut self, colI: usize, rowI: usize) {
         if rowI == self.in_col[colI] { return; }
 
@@ -83,6 +127,8 @@ impl Puzzle {
         self.update(colI, 1);
     }
 
+    /* Algorithm */
+
     fn has_conflicts(&self) -> bool {
         for i in 0..self.N {
             if self.per_row[i] > 1 { return true; }
@@ -93,51 +139,51 @@ impl Puzzle {
         return false;
     }
 
-    fn init(&mut self) {
-        // Clear
-        for i in 0..2 * self.N - 1 {
-            self.per_d1[i] = 0;
-            self.per_d2[i] = 0;
-            if i < self.N {
-                self.per_row[i] = 0;
-            }
-        }
+    fn init_horse(&mut self) {
+        self.clear();
 
-        for col in 0..self.N {
-            let row = self.min_col_row(col);
+        let (mut col, mut row) = (1, 0);
+        for _ in 0..self.N / 2 {
             self.in_col[col] = row;
             self.update(col, 1);
+            col += 2;
+            row += 1;
         }
-        return;
+        let (mut col, mut row) = (0, self.N / 2);
+        for _ in 0..self.N / 2 {
+            self.in_col[col] = row;
+            self.update(col, 1);
+            col += 2;
+            row += 1;
+        }
+    }
 
-        // Horse
-        // let (mut col, mut row) = (1, 0);
-        // for _ in 0..self.N / 2 {
-        //     self.in_col[col] = row;
-        //     self.update(col, 1);
-        //     col += 2;
-        //     row += 1;
-        // }
-        // let (mut col, mut row) = (0, self.N / 2);
-        // for _ in 0..self.N / 2 {
-        //     self.in_col[col] = row;
-        //     self.update(col, 1);
-        //     col += 2;
-        //     row += 1;
-        // }
+    fn init_rand(&mut self) {
+        self.clear();
+
+        self.in_col[0] = self.rng.gen_range(0..self.N);
+        self.update(0, 1);
+
+        for col in 1..self.N {
+            self.in_col[col] = self.min_col_row(col);
+            self.update(col, 1);
+        }
     }
 
     pub fn solve(&mut self) {
+        self.init_horse();
         loop {
-            self.init();
-            for _ in 0..self.N {
+            for _ in 0..2 * self.N {
                 let col = self.max_queen_col();
+                if col == self.N { return; }
+
                 let row = self.min_col_row(col);
                 self.move_queen(col, row);
             }
             if !self.has_conflicts() {
                 break;
             }
+            self.init_rand();
         }
     }
 }
@@ -164,7 +210,7 @@ impl Display for Puzzle {
 
 fn main() {
     let N: usize = read!();
-    if N < 4 {
+    if N == 2 || N == 3 {
         println!("-1");
         return;
     }