example.txt

@@ -1,6 +1,10 @@
-???.### 1,1,3
+...#......
-.??..??...?##. 1,1,3
+.......#..
-?#?#?#?#?#?#?#? 1,3,1,6
+#.........
-????.#...#... 4,1,1
+..........
-????.######..#####. 1,6,5
+......#...
-?###???????? 3,2,1
+.#........
+.........#
+..........
+.......#..
+#...#.....

src/main.rs

@@ -1,118 +1,65 @@
-use std::{collections::HashMap, io::stdin};
+use std::{cmp::max, collections::HashSet, io::stdin};
 
-use anyhow::Result;
+use anyhow::{bail, Result};
-
-fn matches(conditions: &str, group: &str) -> bool {
-    for (g, c) in group.chars().zip(conditions.chars()) {
-        if c == '?' {
-            continue;
-        }
-        if g == c {
-            continue;
-        }
-        return false;
-    }
-    true
-}
-
-type Cache = HashMap<String, i64>;
-
-fn arrangements_inner(conditions: &str, groups: &[String], cache: &mut Cache) -> i64 {
-    let g = &groups[0];
-    let w = conditions.len() - groups.iter().skip(1).map(|s| s.len()).sum::<usize>();
-
-    let positions = w - g.len() + 1;
-    let mut sum = 0;
-    for i in 0..positions {
-        if !matches(&conditions[i..], g) {
-            continue;
-        }
-        if conditions[..i].contains('#') {
-            continue;
-        }
-        let after = if (i + g.len()) < conditions.len() {
-            conditions.chars().nth(i + g.len()).unwrap()
-        } else {
-            '.'
-        };
-        if after == '#' {
-            continue;
-        }
-        if groups.len() == 1 {
-            if conditions[i + g.len()..].contains('#') {
-                continue;
-            }
-            sum += 1;
-        } else {
-            let cache_key = format!("{},{:?}", &conditions[i + g.len()..], &groups[1..]);
-            if let Some(v) = cache.get(&cache_key) {
-                sum += v;
-            } else {
-                let a = arrangements_inner(&conditions[i + g.len()..], &groups[1..], cache);
-                cache.insert(cache_key, a);
-                sum += a;
-            }
-        }
-    }
-    sum
-}
-
-fn arrangements(conditions: &str, groups: &[usize]) -> i64 {
-    let mut group_strings = vec![];
-    group_strings.push("#".repeat(groups[0]).to_owned());
-    for g in groups.iter().skip(1) {
-        group_strings.push(".".to_owned() + &"#".to_owned().repeat(*g));
-    }
-    let mut cache: Cache = Cache::new();
-    arrangements_inner(conditions, &group_strings, &mut cache)
-}
 
 fn main() -> Result<()> {
-    let mut sum = 0;
+    let mut grid = vec![];
-    for line in stdin().lines().map_while(Result::ok) {
+    let mut galaxies = HashSet::new();
-        let mut split = line.split_whitespace();
+    let mut occupied_columns = HashSet::new();
-        let conditions = split.next().unwrap();
+    let mut occupied_rows = HashSet::new();
-        let groups: Vec<usize> = split
+    let mut max_x = 0;
-            .next()
+    let mut max_y = 0;
-            .unwrap()
+    for (y, line) in stdin().lines().map_while(Result::ok).enumerate() {
-            .split(',')
+        let cells: Vec<char> = line.chars().collect();
-            .map(str::parse)
+        max_y = max(max_y, y);
-            .map_while(Result::ok)
+        for (x, c) in cells.iter().enumerate() {
-            .collect();
+            max_x = max(max_x, x);
-
+            if *c == '#' {
-        let conditions = [conditions].repeat(5).join("?");
+                galaxies.insert((x, y));
-        let groups = groups.repeat(5);
+                occupied_columns.insert(x);
-
+                occupied_rows.insert(y);
-        let a = arrangements(&conditions, &groups);
-        sum += a;
             }
+        }
+        grid.push(cells);
+    }
+
+    let mut empty_columns: Vec<usize> = HashSet::from_iter(0..max_x)
+        .difference(&occupied_columns)
+        .copied()
+        .collect();
+    let mut empty_rows: Vec<usize> = HashSet::from_iter(0..max_y)
+        .difference(&occupied_rows)
+        .copied()
+        .collect();
+    empty_columns.sort();
+    empty_rows.sort();
+
+    let mut expanded_galaxies = HashSet::new();
+    let expansion = 1_000_000;
+    for (x, y) in galaxies.into_iter() {
+        let Err(col) = empty_columns.binary_search(&x) else {
+            bail!("x={x} found in empty columns");
+        };
+        let new_x = x + col * (expansion - 1);
+        let Err(row) = empty_rows.binary_search(&y) else {
+            bail!("y={y} found in empty rows");
+        };
+        let new_y = y + row * (expansion - 1);
+        expanded_galaxies.insert((new_x, new_y));
+    }
+    let mut sum = 0;
+    for a in expanded_galaxies.iter() {
+        for b in expanded_galaxies.iter() {
+            let dx = i64::abs(a.0 as i64 - b.0 as i64);
+            let dy = i64::abs(a.1 as i64 - b.1 as i64);
+            // println!("{a:?} -> {b:?} = {}", dx + dy);
+            sum += dx + dy;
+        }
+    }
+    sum /= 2;
     println!("{sum}");
     Ok(())
 }
 
 #[cfg(test)]
-mod tests {
+mod tests {}
-    use super::*;
-
-    #[test]
-    fn test_arrangements() {
-        assert_eq!(arrangements(".??.?#??##???.", &[1, 6]), 5);
-        assert_eq!(arrangements("???.###", &[1, 1, 3]), 1);
-        assert_eq!(arrangements(".??..??...?##.", &[1, 1, 3]), 4);
-        assert_eq!(arrangements("?#?#?#?#?#?#?#?", &[1, 3, 1, 6]), 1);
-        assert_eq!(arrangements("????.#...#...", &[4, 1, 1]), 1);
-        assert_eq!(arrangements("????.######..#####.", &[1, 6, 5]), 4);
-        assert_eq!(arrangements("?###????????", &[3, 2, 1]), 10);
-        assert_eq!(arrangements("????#?#?..?#?", &[4, 1]), 2);
-    }
-
-    #[test]
-    fn test_matches() {
-        assert!(matches("?", "#"));
-        assert!(matches("?", "."));
-        assert!(matches("#", "#"));
-        assert!(matches(".", "."));
-        assert!(!matches("#", "."));
-        assert!(!matches(".", "#"));
-    }
-}