1 files changed, 33 insertions, 8 deletions

diff --git a/src/bin/day8.rs b/src/bin/day8.rs
index ab8595f..8120230 100644
--- a/src/bin/day8.rs
+++ b/src/bin/day8.rs
@@ -4,11 +4,11 @@ static DAY: u8 = 8;
 
 fn main() {
     let input = advent::read_lines(DAY);
-    println!("{DAY}a: {}", Map::new(&input).count_antinodes());
-    println!("{DAY}b: {}", 0);
+    println!("{DAY}a: {}", Map::new(&input).count_antinodes(false));
+    println!("{DAY}b: {}", Map::new(&input).count_antinodes(true));
 }
 
-#[derive(Eq, PartialEq, Hash)]
+#[derive(Eq, PartialEq, Hash, Clone, Copy)]
 struct Position {
     x: isize,
     y: isize,
@@ -45,14 +45,38 @@ impl Map {
         pos.x >= 0 && pos.y >= 0 && pos.x < self.dimensions.x && pos.y < self.dimensions.y
     }
 
-    fn count_antinodes(&self) -> usize {
+    fn get_harmonics(&self, pos1: Position, pos2: Position, with_harmonics: bool) -> HashSet<Position> {
+        let mut harmonics = HashSet::new();
+
+        if with_harmonics {
+            harmonics.insert(pos1);
+            harmonics.insert(pos2);
+        }
+        let mut prev = pos1;
+        let mut harmonic = pos2;
+        loop {
+            (prev, harmonic) = (harmonic, harmonic.add_position(&prev));
+            if !self.inside_map(&harmonic) {
+                break;
+            }
+            harmonics.insert(harmonic);
+            if !with_harmonics {
+                // only interested in first harmonic
+                break;
+            }
+        }
+
+        harmonics
+    }
+
+    fn count_antinodes(&self, with_harmonics: bool) -> usize {
         let mut antinodes = HashSet::new();
 
         for (pos1, antenna) in &self.antennas {
             for (pos2, _) in self.antennas.iter().filter(|&(_, a)| a == antenna).filter(|&(p, _)| p != pos1) {
-                let antinode_pos = pos1.add_position(pos2);
-                if self.inside_map(&antinode_pos) {
-                    antinodes.insert(antinode_pos);
+                let harmonics = self.get_harmonics(*pos1, *pos2, with_harmonics);
+                for antinode in harmonics {
+                    antinodes.insert(antinode);
                 }
             }
         }
@@ -81,6 +105,7 @@ mod tests {
             "............",
             "............",
         ].iter().map(|&x| String::from(x)).collect::<Vec<_>>();
-        assert_eq!(Map::new(&input).count_antinodes(), 14);
+        assert_eq!(Map::new(&input).count_antinodes(false), 14);
+        assert_eq!(Map::new(&input).count_antinodes(true), 34);
     }
 }