summaryrefslogtreecommitdiff
path: root/2020/rust/day07/src/main.rs
blob: d643241fa865debbabff6cf7546161ae969ca39f (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
use anyhow::{anyhow, Result};
use daggy::{Dag, Walker, NodeIndex};
use petgraph::dot::{Dot, Config};
use regex::Regex;
use std::collections::{HashMap, HashSet};
use std::env::args;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter};
use std::io::prelude::*;
use std::path::Path;

type BagCount = usize;
type BagColor = String;
type Rules = Dag<BagColor, BagCount>;
type Index = HashMap<String, NodeIndex>;

// The example data in
// simplified form:    A       | E       |       D
//                       1 x B |   1 x G |      / \
// A contain 1B, 2C.     2 x C |   2 x H |     | A |
// D contain 3B, 4C.   D       | G       |     |/ \|
// B contain 1E.         3 x B |   3 x F |     B __C
// C contain 2E, 9F.     4 x C |   4 x I |     |/  |
// E contain 1G, 2H.   B       | H       |     E   |
// G contain 3F, 4I.     1 x E |   5 x F |    /|  /
// H contain 5F, 6I.   C       |   6 x I |   H G /
// F contain NONE.       2 x E | F       |  / \|/
// I contain NONE.       9 x F | I       | I   F

fn read_input<T: AsRef<Path>>(filename: T) -> Result<(Rules, Index)> {
    let f = File::open(filename)?;
    let reader = BufReader::new(f);

    let outer_re = Regex::new(
        r"(?x)
        ^(?P<outer>[a-z\s]+)\ bags?\ contain\ 
        (?P<inner>(?:(?:\d+\ [a-z\s]+|no\ other)\ bags?[.,]\s?)+)$"
    )?;

    let inner_re = Regex::new(
        r"(?x)
        (?:(?:(?P<count>\d+)\ (?P<color>[a-z\s]+)|(?P<empty>no\ other))\ bags?[.,]\s?)"
    )?;

    let mut rules: Rules = Dag::new();
    let mut rules_index = HashMap::new();

    for line in reader.lines() {
        let line = line?;
        let caps_outer = outer_re.captures(&line).ok_or_else(|| anyhow!("Invalid input"))?;
        let container_color = caps_outer.name("outer")
            .ok_or_else(|| anyhow!("Invalid input"))?.as_str().to_string();
        let parent_index = match rules_index.get(&String::from(&container_color)) {
            None => {
                rules.add_node(container_color.clone())
            },
            Some(existing_index) => {
                *existing_index
            },
        };
        rules_index.insert(container_color, parent_index);
        for m in inner_re.find_iter(caps_outer.name("inner")
            .ok_or_else(|| anyhow!("Invalid input"))?.as_str())
        {
            let s = m.as_str();
            let caps_inner = inner_re.captures(&s).ok_or_else(|| anyhow!("Invalid input"))?;
            match caps_inner.name("empty") {
                Some(_) => { },
                None => {
                    let content_color = caps_inner.name("color")
                        .ok_or_else(|| anyhow!("Invalid input"))?.as_str();
                    let content_count: usize = caps_inner.name("count")
                        .ok_or_else(|| anyhow!("Invalid input"))?.as_str().parse()?;
                    match rules_index.get(content_color) {
                        None => {
                            let (_, child_index) = rules.add_child(
                                parent_index, content_count, String::from(content_color)
                            );
                            rules_index.insert(String::from(content_color), child_index);
                        },
                        Some(existing_index) => {
                            let _ = rules.add_edge(parent_index, *existing_index, content_count);
                        },
                    }
                }
            }
        }
    }

     Ok((rules, rules_index))
}

fn plot_graph(rules: Rules) -> Result<()>{
    let f = File::create("graph.dot")?;
    let mut writer = BufWriter::new(f);
    let g = rules.graph();
    write!(writer, "{}", Dot::with_config(g, &[Config::EdgeNoLabel]))?;
    Ok(())
}

fn recurse_part1(start: NodeIndex, rules: &Rules, mut visited: &mut HashSet<BagColor>) -> usize {
    let mut ans = 0;
    let node = rules.node_weight(start).unwrap();

    if ! visited.contains(node) {
        ans += 1;
        visited.insert(String::from(node));
        let mut parent_walker = rules.parents(start);
        while let Some(walk) = parent_walker.walk_next(&rules) {
            let ( _edge, next_node_index ) = walk;
            ans += recurse_part1(next_node_index, rules, &mut visited);
        }
    }

    ans
}

fn recurse_part2(start: NodeIndex, rules: &Rules) -> usize {
    let mut ans=0;

    rules.node_weight(start).unwrap();
        let mut child_walker = rules.children(start);
        while let Some(walk) = child_walker.walk_next(&rules) {
            let ( edge_index, next_node_index ) = walk;
            let content_count = rules.edge_weight(edge_index).unwrap();
            ans += content_count * (1 + recurse_part2(next_node_index, rules));
        }

    ans
}

fn part1(rules: &Rules, index: &Index) -> Result<usize> {
    let start_index = index.get(&String::from("shiny gold")).ok_or(anyhow!("Bag is missing"))?;
    let mut visited: HashSet<BagColor> = HashSet::new();
    let ans = recurse_part1(*start_index, rules, &mut visited);

    Ok(ans-1)
}

fn part2(rules: &Rules, index: &Index) -> Result<usize> {
    let start_index = index.get(&String::from("shiny gold")).ok_or(anyhow!("Bag is missing"))?;
    Ok(recurse_part2(*start_index, rules))
}

fn main() {
    let ( do_part_1, do_part_2 ) = aoc::do_parts();

    let filename = match args().nth(1) {
        Some(f) => f,
        None => {
            eprintln!("Missing input filename");
            std::process::exit(1);
        },
    };
    match read_input(filename) {
        Ok((input, index)) => {
            if do_part_1 {
                let solution = part1(&input, &index);
                println!("Part1, {:?}", solution);
            }
            if do_part_2 {
                let solution = part2(&input, &index);
                println!("Part2, {:?}", solution);
            }
            match plot_graph(input) {
                Ok(_) => println!("Wrote a graphviz file. Please feel free to plot it."),
                Err(err) => eprintln!("Could not write graphviz file: {}", err),
            }
        },
        Err(err) => eprintln!("Could not read input: {}", err),
    }
}