Day 11: Cosmic Expansion

# Advent of Code 2023
# Day 11: Cosmic Expansion

from itertools import combinations

with open("inputs/day_11.txt") as file:
    image = [list(line[:-1]) for line in file.readlines()]

galaxies = []

for y, row in enumerate(image):
    for x, tile in enumerate(row):
        if tile == "#":
            galaxies.append([x, y])

moved_galaxies_part_1 = [galaxy.copy() for galaxy in galaxies]
moved_galaxies_part_2 = [galaxy.copy() for galaxy in galaxies]

for y, row in enumerate(image):
    if "#" not in row:
        for i, galaxy in enumerate(galaxies):
            if galaxy[1] > y:
                moved_galaxies_part_1[i][1] += 1
                moved_galaxies_part_2[i][1] += 999999

for x in range(len(image[0])):
    if not any(image[y][x] == "#" for y in range(len(image))):
        for i, galaxy in enumerate(galaxies):
            if galaxy[0] > x:
                moved_galaxies_part_1[i][0] += 1
                moved_galaxies_part_2[i][0] += 999999

distance_sum_part_1 = 0

for pair in combinations(moved_galaxies_part_1, 2):
    distance_sum_part_1 += (
        abs(pair[1][0] - pair[0][0]) + abs(pair[1][1] - pair[0][1])
    )

print(f"Part 1: {distance_sum_part_1}")

distance_sum_part_2 = 0

for pair in combinations(moved_galaxies_part_2, 2):
    distance_sum_part_2 += (
        abs(pair[1][0] - pair[0][0]) + abs(pair[1][1] - pair[0][1])
    )

print(f"Part 2: {distance_sum_part_2}")

We can start this one with the most obvious implementation, which is:

Going through all the rows & columns, mark off all the columns and rows that have a galaxy (or, equivalently, all the columns/rows that don't).
Create a new map doubling all those rows/columns and place all the galaxies appropriately in the new map.
Run a BFS from each galaxy to each other galaxy.

This is probably good enough to get you through this part, though it might already be chugging a bit. We'll start by noting that a BFS actually isn't necessary. If you play around a bit, you might be able to see that the shortest path between 2 galaxies is to start at one galaxy and then go vertically until you're on the same plane as the other one, and then just go horizontally until you're at the other galaxy.

To capture this with an algorithm, if we took the X and Y positions of the two galaxies then the distance between those two galaxies is just the absolute differences between the X and Y coordinates (in math terms this is how distance works in the Taxicab Geometry).

We could go further here because the mere mention of "X" and "Y" positions might already be telling you things about how to optimize this, but if we do too much then there's nothing to write for Part 2! So we'll leave it here for now.

This is where the solution above gets quality-checked out of existence. Even if you can check each galaxy in constant time, building a map that's millions by millions doesn't quite work.

But this is where some extra insight into what we're actually doing with our algorithm hints at a way to compress this map. Namely, we don't care about most of the map, we only care about the galaxies!

If we built compressed map that only contains the X and Y positions of the galaxies, then we can just use taxicab distance for all the pairs and solve everything right then and there.

The problem is getting those positions. But let's look more carefully how we're building that map. Namely, if on the un-expanded map a galaxy is at row R and column C, every empty row before R contributes a million* expanded empty rows to R in the expanded map, and every empty column before C independently contributes a million* expanded empty rows to C in the expanded map.

^{*Alright it's a million minus one, because the empty row itself is already counted in R (this is easier to see in Part 1 where it's just two rows).}

That is, in the expanded map,

(expanded-R, expanded-C) = (R + (1000000-1)*(# of empty rows before R), C + (1000000-1)*(# of empty columns before C))

Which we can figure out just by counting the number of empty rows before each R and each C.

Now once we've collected the galaxies, we can mark off which rows/columns are empty and collect a record of all the rows/columns that are empty, and build the expanded map using the above algorithm to get the answer.

Galaxies:
(1,2), (5,3), (2,5), (3,7)

X coords:
1,2,3,5
Total X distance: 
(5-1)+(3-1)+(2-1) + (5-2)+(3-2) + (5-3) =
(5+3+2)-(3*1) + (5+3)-(2*2) + (5)-(1*3) = 13
Y coords:
2,3,5,7
Total Y distance:
(7-2)+(5-2)+(3-2) + (7-3)+(5-3) + (7-5) =
(7+5+3)-(3*2) + (7+3)-(2*3) + (7)-(1*5) = 15

Total distance: 28

#include <bits/stdc++.h>
using namespace std;
int main() {
	string s;
	vector<pair<int,int>> galaxies;
	vector<int> row_is_empty,col_is_empty;
	while(cin >> s) {
		col_is_empty.resize(s.size());
		row_is_empty.push_back(0);
		for(int i=0;i<s.size();i++) {
			if(s[i] != '.') {
				col_is_empty[i] = 1;
				row_is_empty.back() = 1;
				galaxies.emplace_back(row_is_empty.size()-1,i);
			}
		}
	}
	for(int i=0;i<row_is_empty.size();i++) {
		row_is_empty[i] ^= 1;
	}
	for(int i=0;i<col_is_empty.size();i++) {
		col_is_empty[i] ^= 1;
	}
	for(int i=1;i<row_is_empty.size();i++) {
		row_is_empty[i] += row_is_empty[i-1];
	}
	for(int i=1;i<col_is_empty.size();i++) {
		col_is_empty[i] += col_is_empty[i-1];
	}
	for(const int EMPTY_FACTOR: {2,1000000}) {
		long long ans = 0;
		for(int i=0;i<galaxies.size();i++) {
			int r = galaxies[i].first;
			int c = galaxies[i].second;
			r += row_is_empty[r]*(EMPTY_FACTOR-1);
			c += col_is_empty[c]*(EMPTY_FACTOR-1);
			for(int j=0;j<i;j++) {
				int u = galaxies[j].first;
				int v = galaxies[j].second;
				u += row_is_empty[u]*(EMPTY_FACTOR-1);
				v += col_is_empty[v]*(EMPTY_FACTOR-1);
				ans += abs(r-u)+abs(c-v);
			}
		}
		cout << ans << '\n';
	}
}

module Main (main) where
import Data.List (transpose)

main :: IO ()
main = do
    input <- readFile "./input.txt"
    putStrLn $ "Part 1: " ++ show (solve1 $ lines input)
    putStrLn $ "Part 2: " ++ show (solve2 $ lines input)

data Coord = Coord Integer Integer

dist :: Coord -> Coord -> Integer
dist (Coord r1 c1) (Coord r2 c2) = abs (r1 - r2) + abs (c1 - c2)

findGalaxies :: [String] -> Integer -> [Coord]
findGalaxies universe expansionSize = expandedResult where
    findEmpty = map fst . filter (all (=='.') . snd) . zip [0..]
    emptyRows = findEmpty universe
    emptyCols = findEmpty $ transpose universe
    rowIndexed = zip universe [0..]
    rcIndexed = concatMap labelVals rowIndexed where
        labelVals (s, r) = zipWith (\v c -> (Coord r c, v)) s [0..]
    result = map fst $ filter ((=='#') . snd) rcIndexed
    expandedResult = map (expand expansionSize emptyRows emptyCols) result

expand :: Integer -> [Integer] -> [Integer] -> Coord -> Coord
expand expansionSize eRows eCols (Coord r c) = Coord r' c' where
    r' = expanded eRows r
    c' = expanded eCols c
    expanded [] i = i
    expanded (x:xs) i
        | x < i = expansionSize + expanded xs i
        | otherwise = i

solve :: [String] -> Integer -> Integer
solve universe n = result where
    galaxies = findGalaxies universe n
    result = distances galaxies where
        distances [] = 0
        distances (x:xs) = sum (map (dist x) xs) + distances xs

solve1 :: [String] -> Integer
solve1 = flip solve 1

solve2 :: [String] -> Integer
solve2 = flip solve 999999

The idea is to store the x/y distances as a Prefix sum. Each row/column on the map is assigned the distance light travels when crossing that row/column, and thanks to the prefix sum I can lookup quickly how far apart two x or two y coordinates are.

And since the distance between two galaxies is the manhattan distance here (because we cannot move diagonally), knowing the distances along the x and y axes is enough to compute the distance between galaxies.

pub struct PuzzleInput {
  [...]
}

impl PuzzleInput {
    pub fn shortest_path(&self, start: (usize, usize), goal: (usize, usize)) -> usize {
        let (x1, y1) = start;
        let (x2, y2) = goal;

        let dx = self.col_dist.query(x1, x2);
        let dy = self.row_dist.query(y1, y2);

        dx + dy
    }
}

[...]

pub struct PrefixSum {
    prefix_sums: Vec<usize>,
}

impl PrefixSum {
    pub fn new(values: &[usize]) -> Self {
        let mut prefix_sums = Vec::with_capacity(values.len());

        for v in values {
            prefix_sums.push(prefix_sums.last().unwrap_or(&0) + v);
        }

        Self { prefix_sums }
    }

    pub fn query(&self, x1: usize, x2: usize) -> usize {
        let (x1, x2) = (x1.min(x2), x1.max(x2));

        let lower = if x1 == 0 { 0 } else { self.prefix_sums[x1 - 1] };
        let upper = if x2 == 0 { 0 } else { self.prefix_sums[x2 - 1] };

        upper - lower
    }
}

day_11_bench  fastest       │ slowest       │ median        │ mean          │ samples │ iters
├─ parse      63.6 µs       │ 204.3 µs      │ 66.98 µs      │ 69.25 µs      │ 100     │ 100
├─ part1      303.6 µs      │ 478.3 µs      │ 310.9 µs      │ 317.7 µs      │ 100     │ 100
╰─ part2      303.2 µs      │ 324.9 µs      │ 307.9 µs      │ 310.1 µs      │ 100     │ 100

I am irritated with myself for not planning more for part 1. I just sort of jumped into implementation and wasted a lot of time expanding the grid array and dealing with the fact that once you
start expanding it, your row numbers are for the old grid, so you have to account for that.

It cost me a lot of time, so I felt pretty dumb when I got into part2 and it made me realize I didn't actually need to expand the grid, just track how many empty rows and columns are crossed on the shortest path.

For computing the distances, I at least noticed that you didn't actually have to stair step down, it is just the sum of the row difference plus the column difference.

I realized my VisitCells function was written with the assumption of a dense grid, so I added a VisitActive function that only visits populated cells, more useful for a sparse grid.

visiting all unique pairs seemed like a useful generic function too, so I added that to the library instead of putting it in the code just for today.

--- a/grid/grid.go
+++ b/grid/grid.go
@@ -9,15 +9,20 @@ type Position struct {
        Col int
 }

+type Pair [2]Position
+
 type Grid[T any] struct {
        Items      map[Position]T
        Height     int
        Width      int
        Wraparound bool
        Diagonal   bool
 }

 type CellVisitor[T any] func(pos Position, value T, ok bool)

+type PairVisitor[T any] func(pos0, pos1 Position, value0, value1 T)
+
 func NewGrid[T any](height int, width int) *Grid[T] {
        g := Grid[T]{
                Items:  make(map[Position]T),
@@ -49,6 +54,12 @@ func (g *Grid[T]) Set(rowIdx int, colIdx int, value T) {
        g.SetAt(pos, value)
 }

+func (g *Grid[T]) VisitActive(visitor CellVisitor[T]) {
+       for pos, value := range g.Items {
+               visitor(pos, value, true)
+       }
+}
+
 func (g *Grid[T]) VisitCells(visitor CellVisitor[T]) {
        for rowIdx := 0; rowIdx < g.Height; rowIdx++ {
                for colIdx := 0; colIdx < g.Width; colIdx++ {
@@ -75,6 +86,26 @@ func (g *Grid[T]) VisitCol(colIdx int, visitor CellVisitor[T]) {
        }
 }
+func (g *Grid[T]) VisitPairs(visitor PairVisitor[T]) {
+       history := make(map[Pair]bool)
+       for pos0, value0 := range g.Items {
+               for pos1, value1 := range g.Items {
+                       if pos0 == pos1 {
+                               continue
+                       }
+
+                       visited := history[Pair{pos1, pos0}]
+                       if visited {
+                               continue
+                       }
+
+                       history[Pair{pos0, pos1}] = true
+                       visitor(pos0, pos1, value0, value1)
+               }
+       }
+}
+
 func (g *Grid[T]) VisitNeighbors(rowIdx int, colIdx int, visitor CellVisitor[T]) {
        for rowDelta := -1; rowDelta <= 1; rowDelta++ {
                for colDelta := -1; colDelta <= 1; colDelta++ {

--- a/grid/grid_test.go
+++ b/grid/grid_test.go
@@ -37,11 +37,23 @@ func TestVisitors(t *testing.T) {

        g := NewGrid[int](10, 10)

+       count := 0
+       g.VisitActive(func(pos Position, value int, ok bool) {
+               count += 1
+       })
+       assert.Equal(count, 0)
+
        g.VisitCells(func(pos Position, value int, ok bool) {
                assert.False(ok)
                g.SetAt(pos, 1)
        })

+       count = 0
+       g.VisitActive(func(pos Position, value int, ok bool) {
+               count += 1
+       })
+       assert.Equal(count, 100)
+
        expectedColIdx := 0
        g.VisitRow(5, func(pos Position, value int, ok bool) {
                assert.Equal(value, 1)
@@ -124,3 +163,19 @@ func TestNeighborsWithWraparound(t *testing.T) {
        assert.Equal(neighbors[6], Position{1, 0})
        assert.Equal(neighbors[7], Position{1, 1})
 }
+
+func TestPairs(t *testing.T) {
+       assert := assert.New(t)
+
+       g := NewGrid[bool](10, 10)
+
+       g.Set(0, 0, true)
+       g.Set(5, 5, true)
+       g.Set(9, 9, true)
+
+       pairCount := 0
+       g.VisitPairs(func(pos0, pos1 Position, value0, value1 bool) {
+               pairCount++
+       })
+       assert.Equal(pairCount, 3)
+}

a useful trick is that the distance between two galaxies is the Taxicab distance and you don't actually need to search for it, it's just the distance in rows plus the distance in columns.

package day11

import (
    "math"

    "github.com/rs/zerolog/log"
    "github.com/spf13/cobra"

    "spit-evil-olive.tips/aoc2023/common"
    "spit-evil-olive.tips/aoc2023/grid"
)

func parseInput(path string) (*grid.Grid[bool], error) {
    lines, err := common.ReadFileLines(path)
    if err != nil {
        return nil, err
    }

    height := len(lines)
    width := len(lines[0])
    universe := grid.NewGrid[bool](height, width)

    for rowIdx, line := range lines {
        for colIdx, value := range line {
            if value == '#' {
                universe.Set(rowIdx, colIdx, true)
            }
        }
    }

    return universe, nil
}

func getExpansionMapping(occupiedMap map[int]bool, size int) (map[int]int, int) {
    mapping := make(map[int]int)

    offset := 0
    for src := 0; src < size; src++ {
        occupied, _ := occupiedMap[src]
        if occupied {
            mapping[src] = src + offset
        } else {
            offset += 1
        }
    }

    newSize := size + offset
    return mapping, newSize
}

// a quantum leap forward
// in time and space
// the universe learned to expand
func expand(universe *grid.Grid[bool]) *grid.Grid[bool] {
    occupiedRows := make(map[int]bool)
    occupiedCols := make(map[int]bool)

    universe.VisitActive(func(pos grid.Position, value bool, ok bool) {
        occupiedRows[pos.Row] = true
        occupiedCols[pos.Col] = true
    })

    rowMapping, newHeight := getExpansionMapping(occupiedRows, universe.Height)
    colMapping, newWidth := getExpansionMapping(occupiedCols, universe.Width)

    newUniverse := grid.NewGrid[bool](newHeight, newWidth)

    universe.VisitActive(func(pos grid.Position, value bool, ok bool) {
        newRow, ok1 := rowMapping[pos.Row]
        newCol, ok2 := colMapping[pos.Col]

        if !ok1 || !ok2 {
            log.Panic().Msgf("failed to map %v", pos)
        }

        newUniverse.Set(newRow, newCol, true)
    })

    return newUniverse
}

var CommandA = &cobra.Command{
    Use:  "11a",
    Args: cobra.ExactArgs(1),
    RunE: func(cmd *cobra.Command, args []string) error {
        universe, err := parseInput(args[0])
        if err != nil {
            return err
        }

        newUniverse := expand(universe)

        sum := 0
        newUniverse.VisitPairs(func(pos0, pos1 grid.Position, value0, value1 bool) {
            rowDist := math.Abs(float64(pos0.Row - pos1.Row))
            colDist := math.Abs(float64(pos0.Col - pos1.Col))
            distance := rowDist + colDist
            sum += int(distance)
        })

        log.Printf("%v", sum)

        return nil
    },
}

the sparse implementation of my Grid ends up being a huge help here.

--- a/day11/day11a.go
+++ b/day11/day11a.go
@@ -40,7 +40,7 @@ func getExpansionMapping(occupiedMap map[int]bool, size int) (map[int]int, int)
                if occupied {
                        mapping[src] = src + offset
                } else {
-                       offset += 1
+                       offset += 999999
                }
        }

> env BENCHMARK=100 go run main.go 11a day11/11.txt
{"level":"info","time":"2023-12-10T22:24:31.310705547-08:00","message":"mean: 24.503669ms"}
{"level":"info","time":"2023-12-10T22:24:31.310772596-08:00","message":"stddev: 521.422µs"}
> env BENCHMARK=100 go run main.go 11a day11/11.txt
{"level":"info","time":"2023-12-10T22:24:16.085930636-08:00","message":"mean: 24.819151ms"}
{"level":"info","time":"2023-12-10T22:24:16.085992656-08:00","message":"stddev: 654.189µs"}

no meaningful difference in runtime between parts A and B, as expected. but running over the example input yields a noticeable difference:

> env BENCHMARK=100 go run main.go 11a day11/11e.txt
{"level":"info","time":"2023-12-10T22:34:41.490831248-08:00","message":"mean: 29.072µs"}
{"level":"info","time":"2023-12-10T22:34:41.490883071-08:00","message":"stddev: 8.455µs"}

140x140 instead of 10x10, so about 200 times larger input yielding about 1000 times slower performance.

I had a suspicion about where my bottleneck might live - I implemented VisitPairs in a very brute-force way. the algorithm is inherently O(N^2) in runtime, but my simple implementation also uses O(N^2) memory, which is suboptimal. here's a better version that only needs O(N) memory, and allocates it in one shot to avoid resizing the array:

func (g *Grid[T]) VisitPairs(visitor PairVisitor[T]) {
    type Item struct {
        pos   Position
        value T
    }

    items := make([]Item, 0, len(g.Items))
    for pos, value := range g.Items {
        items = append(items, Item{pos, value})
    }

    for i := 0; i < len(items); i++ {
        item0 := items[i]
        for j := i+1; j < len(items); j++ {
            item1 := items[j]
            visitor(item0.pos, item1.pos, item0.value, item1.value)
        }
    }
}

and I was expecting that to give me a speedup, but holy crap:

> env BENCHMARK=100 go run main.go 11a day11/11e.txt
{"level":"info","time":"2023-12-10T23:05:36.2297823-08:00","message":"mean: 21.331µs"}
{"level":"info","time":"2023-12-10T23:05:36.229821551-08:00","message":"stddev: 8.96µs"}
> env BENCHMARK=100 go run main.go 11a day11/11.txt
{"level":"info","time":"2023-12-10T23:05:48.420022398-08:00","message":"mean: 455.641µs"}
{"level":"info","time":"2023-12-10T23:05:48.420065281-08:00","message":"stddev: 88.818µs"}

For second part I tried adding the factor variable in, but that only cause OOM. So, I added virtualExpand that recompute star positions based on the rows. While doing that, I totally forgot about the mutating mid-loop thing...

import kotlin.math.abs

fun main() {
//    val factor = 2
    val factor = 1_000_000

    val input = mutableListOf<MutableList<Char>>()
    while (true) {
        input.add(readLine()?.toCharArray()?.toMutableList() ?: break)
    }
    val colsToExpand = findExpandingCol(input)
    val rowsToExpand = findExpandingRow(input)

    // Part 1
//    expandRows(input, rowsToExpand, factor)
//    expandCols(input, colsToExpand, factor)
//    printMap(input)
//    val stars = findStars(input)

    // Part 2
    var stars = findStars(input)
    stars = virtualExpand(stars, rowsToExpand, colsToExpand, factor)
    println(stars)

    val starPairs = combination(stars)
    val sumDistance = starPairs.map { distanceBetween(it.first, it.second) }.sum()
    println(sumDistance)
}

fun virtualExpand(stars: List<Pair<Int, Int>>, rowsToExpand: List<Int>, colsToExpand: List<Int>, factor: Int): List<Pair<Int, Int>> {
    return stars.map {
//        println("expanding $it")
        var row = it.first
        for(expandingRow in rowsToExpand) {
            if (expandingRow < it.first){
                row += factor-1
            }
        }
        var col = it.second
        for(expandingCol in colsToExpand) {
            if (expandingCol < it.second){
//                println("expanding col $expandingCol")
                col += factor-1
            }
        }

//        println("end result ($row, $col)")
        row to col
    }
}

fun printMap(l: List<List<Char>>) {
    for(row in l) {
        for (col in row) {
            print(col)
        }
        println()
    }
}

fun findExpandingCol(l: List<List<Char>>): List<Int> {
    return l[0].indices.filter { col ->
        l.all { it[col] == '.' }
    }
}

fun findExpandingRow(l: List<List<Char>>): List<Int> {
    return l.indices.filter { row ->
        l[row].all { it == '.' }
    }
}

fun findStars(l: List<List<Char>>): List<Pair<Int, Int>> {
    val out = mutableListOf<Pair<Int, Int>>()
    l.mapIndexed { rowId, row ->
        row.mapIndexed { colId, col ->
            if (col == '#') {
                out.add(rowId to colId)
            }
        }
    }
    return out
}

fun expandRows(input: MutableList<MutableList<Char>>, rowsToExpand: List<Int>, factor: Int = 1) {
    var expanded = 0
    for (row in rowsToExpand) {
        val pattern = (0..<input[0].size).map { '.' }.toMutableList()
        val injectedRows = (0..<(factor-1)).map { pattern.toMutableList() }
        val realRow = row + expanded
        input.addAll(realRow, injectedRows)
        expanded += factor
    }
}

fun expandCols(input: MutableList<MutableList<Char>>, colsToExpand: List<Int>, factor: Int = 1) {
    val injectedCols = (0..<(factor-1)).map { '.' }
    for (row in input) {
        var expanded = 0
        for (col in colsToExpand) {
            val realCol = col + expanded
            row.addAll(realCol, injectedCols)
            expanded += factor
        }
    }
}

fun <T> combination(list: List<T>) = sequence {
    list.indices.forEach { a->
        ((a+1)..<list.size).forEach { b->
            yield(list[a] to list[b])
        }
    }
}

fun distanceBetween(first: Pair<Int, Int>, second: Pair<Int, Int>): Long {
    return abs(first.first - second.first).toLong() + abs(first.second - second.second).toLong()
}

I had some extra time to tweak things and see how they affected performance, so parts of the code are a bit more complicated or weird than they need to be, but hopefully it still mostly makes sense.

To avoid having to deal with a sparsely-filled grid after expanding along one axis, I did the expansions separately (and concurrently), with each producing a map looking like original_galaxy_coordinates => new_x_or_y_coordinate. Then, I merged the two maps together to get the full new coordinates of each galaxy.

defmodule AdventOfCode.Solution.Year2023.Day11 do
  alias AdventOfCode.CharGrid, as: G

  def part1(input), do: solve(input, 1)
  def part2(input), do: solve(input, 999_999)

  def solve(input, spacing) do
    input
    |> G.from_input()
    |> expanded_galaxy_coords(spacing)
    |> pairwise_distances_sum()
  end

  defp expanded_galaxy_coords(grid, spacing) do
    [new_xs, new_ys] =
      Task.await_many([
        Task.async(fn -> expanded_axis(grid, spacing, &G.cols/1, fn {x, _y} -> x end) end),
        Task.async(fn -> expanded_axis(grid, spacing, &G.rows/1, fn {_x, y} -> y end) end)
      ])

    Enum.map(new_xs, fn {coords, x} -> {x, new_ys[coords]} end)
  end

  defp expanded_axis(grid, spacing, lanes_fn, axis_fn) do
    grid
    |> lanes_fn.()
    |> Enum.flat_map_reduce(0, fn lane, expand_by ->
      if Enum.all?(lane, &match?({_coords, ?.}, &1)) do
        {[], expand_by + spacing}
      else
        {for({coords, ?#} <- lane, do: {coords, axis_fn.(coords) + expand_by}), expand_by}
      end
    end)
    |> then(fn {new_axis, _expand_by} -> Map.new(new_axis) end)
  end

  defp pairwise_distances_sum(galaxy_coords, sum \\ 0)
  defp pairwise_distances_sum([], sum), do: sum

  defp pairwise_distances_sum([{x1, y1} | rest], sum) do
    new_sums = for({x2, y2} <- rest, reduce: 0, do: (acc -> acc + abs(x2 - x1) + abs(y2 - y1)))
    pairwise_distances_sum(rest, sum + new_sums)
  end
end

use std::env;
use std::io::{self, prelude::*, BufReader};
use std::fs::File;

use point2d::point2d::Point2D;

extern crate itertools;
use itertools::Itertools;

// Expand map, where times parameter is the column multiplier
// eg. expanding the map doubles (times=2) the width of empty space
fn expand_map(map: &Vec<Point2D>, times: i64) -> Vec<Point2D> {
    let xmax = map.iter().map(|pt| pt.x).max().unwrap();
    let ymax = map.iter().map(|pt| pt.y).max().unwrap();

    // Determine which rows and columns to expand
    let expanded_rows: Vec<_> =
        (0..=ymax)
        .into_iter()
        .filter(|y| {
            (0..=xmax).into_iter().all(|x| !map.contains(&Point2D { x: x, y: *y}))
        })
        .collect();
    let expanded_cols: Vec<_> =
        (0..=xmax)
        .into_iter()
        .filter(|x| {
            (0..=ymax).into_iter().all(|y| !map.contains(&Point2D { x: *x, y: y}))
        })
        .collect();

    // Generate expanded map
    let mut new_map = map.clone();
    for y in (0..=ymax).rev() {
        if expanded_rows.contains(&y) {
            for g in &mut new_map {
                if g.y > y { g.y += times-1 }
            }
        }
    }
    for x in (0..=xmax).rev() {
        if expanded_cols.contains(&x) {
            for g in &mut new_map {
                if g.x > x { g.x += times-1 }
            }
        }
    }
    new_map
}

fn distance(one: &Point2D, other: &Point2D) -> i64 {
    (one.x - other.x).abs() + (one.y - other.y).abs()
}

fn distance_sum(galaxies: &Vec<Point2D>) -> i64 {
    galaxies
        .iter()
        .combinations(2)
        .map(|g| distance(g[0],g[1]))
        .sum::<i64>()
}

fn solve(input: &str) -> io::Result<()> {
    let file = File::open(input).expect("Input file not found.");
    let reader = BufReader::new(file);

    // Input
    let input: Vec<String> = match reader.lines().collect() {
        Err(err) => panic!("Unknown error reading input: {err}"),
        Ok(result) => result,
    };

    // Build map
    let mut galaxies: Vec<Point2D> = Vec::new();
    for (y,line) in input.iter().enumerate() {
        for (x,ch) in line.chars().enumerate() {
            let pt = Point2D { x: x as i64, y: y as i64 };
            if ch == '#' {
                galaxies.push(pt);
            }
        }
    }

    // Part 1 + Part 2
    let part1 = distance_sum(&expand_map(&galaxies,2));
    let part2 = distance_sum(&expand_map(&galaxies,1_000_000));
    println!("Part 1: {part1}"); // 9509330
    println!("Part 2: {part2}"); //635832237682

    Ok(())
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let filename = &args[1];
    solve(&filename).unwrap();
}