1 files changed, 30 insertions, 22 deletions

diff --git a/src/structs.rs b/src/structs.rs
index fdc20f4..aa70680 100644
--- a/src/structs.rs
+++ b/src/structs.rs
@@ -9,27 +9,28 @@ impl StoneWall {
         let mut rows = vec![vec![0; n]; h];
         rows.get_mut(0).map(|r| r[0] = 1);
         rows.get_mut(1).map(|r| r[1] = 2);
-        Self {
-            rows,
-        }
+        Self { rows }
     }
 
     pub fn set_stone(&mut self, row: u32, pos: u32, stone: u32) -> Option<()> {
-        self.rows.get_mut(row as usize).and_then(|v| v.get_mut(pos as usize))
+        self.rows
+            .get_mut(row as usize)
+            .and_then(|v| v.get_mut(pos as usize))
             .map(|v| *v = stone)
     }
 
-    pub fn output(&self) {
-        let colors = [
-            [31, 32], 
-            [33, 35],
-        ];
+    pub fn output(&mut self) {
+        let colors = [[31, 32], [33, 35]];
+        self.rows.sort_by_key(|x| x[0]);
+
         for (i, row) in self.rows.iter().enumerate() {
             for (j, &stone) in row.iter().enumerate() {
-                print!("{}", colors[i & 1][j & 1]);
-                print!("\x1b[{}m{}", 
-                       colors[i & 1][j & 1],
-                       "◙".repeat(stone as usize));
+                //print!("{}", colors[i & 1][j & 1]);
+                print!(
+                    "\x1b[{}m{}",
+                    colors[i & 1][j & 1],
+                    "◙".repeat(stone as usize)
+                );
             }
             println!("\x1b[m");
         }
@@ -80,14 +81,23 @@ impl GapHeights {
     pub fn output(&self, n: u32, h: u32) {
         let mut stones = vec![0; n as usize];
         let mut toggle = 0;
-        let colors = [
+        let mut colors = [
             "\x1b[31m", "\x1b[32m", "\x1b[33m", "\x1b[34m", "\x1b[35m", "\x1b[36m",
-        ];
-        for row in 0..h {
+        ]
+        .iter()
+        .cycle();
+        let mut indices = Vec::with_capacity(h as usize);
+        for i in 0..n {
+            let n = self.heights[i as usize];
+            if indices.iter().all(|x| *x != n) {
+                indices.push(n);
+            }
+        }
+        //println!("{:?} : {:?}", indices, self.heights);
+        for row in indices {
             self.calculate_row(row, &mut stones);
             for &len in stones.iter() {
-                print!("{}", colors[toggle]);
-                toggle = (toggle + 1) % colors.len();
+                print!("{}", colors.next().unwrap());
                 for _ in 0..len {
                     print!("◙");
                 }
@@ -104,16 +114,14 @@ impl GapHeights {
     }
 
     pub fn as_stone_wall(&self, n: u32) -> StoneWall {
-        let h = n/2 + 1;
+        let h = n / 2 + 1;
         let mut rows = Vec::with_capacity(h as usize);
         for i in 0..h {
             let mut row = vec![0; n as usize];
             self.calculate_row(i, &mut row);
             rows.push(row);
         }
-        StoneWall {
-            rows
-        }
+        StoneWall { rows }
     }
 }