advent-of-code/2018/day17/src/main.rs

use std::collections::VecDeque;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::ops::Range;

const MAP_SIZE: usize = 1800;
const FOUNTAIN: char = '+';
const EMPTY_SPACE: char = ' ';
const WALL_SPACE: char = '#';
const WATER: char = '~';
const FALLING_WATER: char = '|';
type Map = Vec<Vec<char>>;

#[derive(Debug, PartialEq)]
struct Coord {
    x: usize,
    y: usize,
}

struct MapBounds {
    min_y: usize,
    max_y: usize,
}

fn find_range_in_line(line: &str) -> (Range<usize>, Range<usize>) {
    let mut parts = line.split(", ");
    let mut left_parts = parts.nth(0).unwrap().split("=");
    let left_coord_type: &str = left_parts.nth(0).unwrap();
    let left_number: usize = left_parts.nth(0).unwrap().parse().unwrap();
    let mut right_parts = parts.nth(0).unwrap().split("=");
    let mut right_numbers = right_parts.nth(1).unwrap().split("..");
    let right_number_1: usize = right_numbers.nth(0).unwrap().parse().unwrap();
    let right_number_2: usize = right_numbers.nth(0).unwrap().parse().unwrap();
    let left_range = Range { start: left_number, end: left_number + 1 };
    let right_range = Range { start: right_number_1, end: right_number_2 + 1 };
    if left_coord_type == "x" {
        (left_range, right_range)
    } else {
        (right_range, left_range)
    }
}

fn print_map(map: &Map, min_x: usize, max_x: usize, min_y: usize, max_y: usize) {
    for y in min_y..=max_y {
        for x in min_x..=max_x {
            print!("{}", map[x][y]);
        }
        println!();
    }
}

fn build_map_from_file(filename: &str) -> (Map, MapBounds) {
    let mut min_y: usize = MAP_SIZE;
    let mut max_y: usize = 0;
    let mut m = Vec::<Vec<char>>::with_capacity(MAP_SIZE);
    for _ in 0..MAP_SIZE {
        let mut n = Vec::<char>::with_capacity(MAP_SIZE);
        for _ in 0..MAP_SIZE {
            n.push(EMPTY_SPACE);
        }
        m.push(n);
    }
    let lines: Vec<String> = BufReader::new(File::open(filename).unwrap())
        .lines()
        .map(|line| line.unwrap())
        .collect();
    m[500][0] = FOUNTAIN;
    for line in lines.iter() {
        let (x_range, y_range) = find_range_in_line(line.as_str());
        for x in x_range {
            for y in y_range.clone() {
                if y < min_y {
                    min_y = y;
                }
                if y > max_y {
                    max_y = y;
                }
                m[x][y] = '#';
            }
        }
    }
    (m, MapBounds { min_y, max_y })
}

fn pour_over(m: &mut Map, start_x: usize, start_y: usize) {
    let mut frontier: VecDeque<Coord> = VecDeque::new();
    frontier.push_back(Coord { x: start_x, y: start_y });

    while let Some(coord) = frontier.pop_front() {
        match m[coord.x][coord.y] {
            FOUNTAIN => frontier.push_back(Coord { x: coord.x, y: coord.y + 1 }),
            EMPTY_SPACE => {
                m[coord.x][coord.y] = FALLING_WATER;
                if coord.y + 1 < MAP_SIZE {
                    match m[coord.x][coord.y + 1] {
                        WALL_SPACE => frontier.push_back(Coord { x: coord.x, y: coord.y }),
                        EMPTY_SPACE => frontier.push_back(Coord { x: coord.x, y: coord.y + 1 }),
                        FALLING_WATER => frontier.push_back(Coord { x: coord.x, y: coord.y }),
                        WATER => {
                            if m[coord.x + 1][coord.y] == WALL_SPACE || m[coord.x - 1][coord.y] == WALL_SPACE {
                                frontier.push_back(Coord { x: coord.x, y: coord.y });
                            } else if m[coord.x + 2][coord.y] == WALL_SPACE || m[coord.x - 2][coord.y] == WALL_SPACE {
                                frontier.push_back(Coord { x: coord.x, y: coord.y });
                            } else if m[coord.x + 1][coord.y] == FALLING_WATER || m[coord.x - 1][coord.y] == FALLING_WATER {
                                frontier.push_back(Coord { x: coord.x, y: coord.y });
                            }
                        },
                        _ => {},
                    }
                }
            },
            WALL_SPACE => {},
            FALLING_WATER => {
                let mut found_down = false;
                // go left as far as possible
                for x in (0..coord.x).rev() {
                    match m[x][coord.y] {
                        EMPTY_SPACE => {
                            if m[x][coord.y + 1] == EMPTY_SPACE {
                                found_down = true;
                                let c = Coord { x: x, y: coord.y };
                                if !frontier.contains(&c) {
                                    frontier.push_back(c);
                                }
                                break;
                            } else {
                                m[x][coord.y] = WATER;
                            }
                        },
                        WATER => {},
                        FALLING_WATER => {},
                        _ => break,
                    }
                }
                // go right as far as possible
                for x in (coord.x + 1)..MAP_SIZE {
                    match m[x][coord.y] {
                        EMPTY_SPACE => {
                            if m[x][coord.y + 1] == EMPTY_SPACE {
                                found_down = true;
                                let c = Coord { x: x, y: coord.y };
                                if !frontier.contains(&c) {
                                    frontier.push_back(c);
                                }
                                break;
                            } else {
                                m[x][coord.y] = WATER;
                            }
                        },
                        WATER => {},
                        FALLING_WATER => {},
                        _ => break,
                    }
                }
                // if no down, then push the square up
                if !found_down {
                    let c = Coord { x: coord.x, y: coord.y - 1 };
                    if !frontier.contains(&c) {
                        frontier.push_back(c);
                    }
                }
            },
            WATER => {},
            _ => panic!("invalid map square at {:#?}", coord),
        }
    }
}

fn lock_water(m: &mut Map, bounds: &MapBounds) {
    for y in bounds.min_y..=bounds.max_y {
        for x in 0..MAP_SIZE {
            if m[x][y] == FALLING_WATER {
                let mut locked = true;
                for k in (0..x).rev() {
                    if m[k][y] == WALL_SPACE {
                        break;
                    } else if m[k][y] == EMPTY_SPACE {
                        locked = false;
                        break;
                    }
                }
                for k in (x+1)..MAP_SIZE {
                    if m[k][y] == WALL_SPACE {
                        break;
                    } else if m[k][y] == EMPTY_SPACE {
                        locked = false;
                        break;
                    }
                }
                if locked {
                    m[x][y] = WATER;
                } else {
                    for k in (0..x).rev() {
                        if m[k][y] == EMPTY_SPACE || m[k][y] == WALL_SPACE {
                            break;
                        }
                        m[k][y] = FALLING_WATER;
                    }
                    for k in (x+1)..MAP_SIZE {
                        if m[k][y] == EMPTY_SPACE || m[k][y] == WALL_SPACE {
                            break;
                        }
                        m[k][y] = FALLING_WATER;
                    }
                }
            }
        }
    }
}

fn count_water(m: &Map, bounds: &MapBounds) -> usize {
    let mut squares: usize = 0;
    for y in bounds.min_y..=bounds.max_y {
        for x in 0..MAP_SIZE {
            if m[x][y] == WATER || m[x][y] == FALLING_WATER {
                squares += 1;
            }
        }
    }
    squares
}

fn count_resting_water(m: &Map, bounds: &MapBounds) -> usize {
    let mut squares: usize = 0;
    for y in bounds.min_y..=bounds.max_y {
        for x in 0..MAP_SIZE {
            if m[x][y] == WATER {
                squares += 1;
            }
            if m[x][y] == FALLING_WATER {
                if m[x - 1][y] == WATER || m[x + 1][y] == WATER {
                    squares += 1;
                }
            }
        }
    }
    squares
}

fn main() {
    let (mut m, bounds) = build_map_from_file("input");
    pour_over(&mut m, 500, 0);
    if cfg!(debug_assertions) {
        print_map(&m, 0, MAP_SIZE - 1, 0, MAP_SIZE - 1);
    }
    let water_squares = count_water(&m, &bounds);
    println!("found {} water squares", water_squares);
    lock_water(&mut m, &bounds);
    if cfg!(debug_assertions) {
        print_map(&m, 0, MAP_SIZE - 1, 0, MAP_SIZE - 1);
    }
    let water_squares = count_water(&m, &bounds);
    println!("found {} locked water squares", water_squares);
    let resting_squares = count_resting_water(&m, &bounds);
    println!("found {} resting water squares", resting_squares);
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_range_from_line() {
        let line = "x=100, y=100..102";
        let (x_range, y_range) = find_range_in_line(line);
        assert_eq!(x_range, (100..101));
        assert_eq!(y_range, (100..103));
    }

    #[test]
    fn test_backwards_range_from_line() {
        let line = "y=99, x=1..5";
        let (x_range, y_range) = find_range_in_line(line);
        assert_eq!(x_range, (1..6));
        assert_eq!(y_range, (99..100));
    }

    #[test]
    fn test_sample_data() {
        let (mut m, bounds) = build_map_from_file("test-input");
        pour_over(&mut m, 500, 0);
        let water_squares = count_water(&m, &bounds);
        assert_eq!(water_squares, 57);
        let resting_squares = count_resting_water(&m, &bounds);
        assert_eq!(resting_squares, 29);
    }
}