Dec 21, 2020

12 min #rust · #advent-of-code

Day 7 (Advent of Code 2020)

From the series Advent of Code 2020

14

👋 This page was last updated ~4 years ago. Just so you know.

Another day, another Advent of Code 2020 problem.

That one seems fun! For some nerdy values of fun.

Our input is a set of rules:

light red bags contain 1 bright white bag, 2 muted yellow bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
bright white bags contain 1 shiny gold bag.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
faded blue bags contain no other bags.
dotted black bags contain no other bags.

And the question we have to answer is: “How many bag colors can eventually contain at least one shiny gold bag?”

From looking at the example input above, for example, “bright white bags” can directly contain a shiny gold bag, whereas light red bags can contain a bright white bag which can contain a shiny gold bag.

The question for part 1 doesn’t use “how many bags of colors Y can be contained in bags of color X”, so I suspect it’s going to show up in part 2 (which, as a reminder, I can’t see until I’ve solved the first part) - I’d like to try and parse and represent that input as types that make it easy to answer that first question.

Bags are clearly defined by two things: an adjective (light, dark, bright, muted, shiny, vibrant, faded, dotted) and a color (red, orange, white, yellow, gold, olive, plum, etc).

We don’t know yet how many different adjectives or colors there’ll be, so I’m tempted to just represent those as borrowed strings:

/// (adjective, color), i.e. ("dark", "orange")
type BagSpec<'a> = (&'a str, &'a str);

Cool Bear's Hot Tip

“Spec” is used here as a shorthand for “specification”.

From there, all our rules are simply a mapping from BagSpec to a quantity and another BagSpec:

use std::collections::HashMap;

/// K can contain V.0 of V.1
type Rules<'a> = HashMap<BagSpec<'a>, (usize, BagSpec<'a>)>;

But here’s the thing - some bags can contain several types of other bags. From the example above:

light red bags contain 1 bright white bag, 2 muted yellow bags.

We can’t represent this rule with our current set of types:

fn main() {
    let mut rules: Rules = Default::default();
    rules.insert(("light", "red"), (1, ("bright", "white")));
    rules.insert(("light", "red"), (2, ("muted", "yellow")));
    dbg!(&rules);
}

$ cargo run --quiet
[src/main.rs:13] &rules = {
    (
        "light",
        "red",
    ): (
        2,
        (
            "muted",
            "yellow",
        ),

The second rule has overwritten the first one!

The good news is: there’s a crate (and a type) for that!

$ cargo add multimap
      Adding multimap v0.8.2 to dependencies

use multimap::MultiMap;

/// K can contain V.0 of V.1
type Rules<'a> = MultiMap<BagSpec<'a>, (usize, BagSpec<'a>)>;

$ cargo run --quiet
[src/main.rs:13] &rules = {
    (
        "light",
        "red",
    ): [
        (
            1,
            (
                "bright",
                "white",
            ),
        ),
        (
            2,
            (
                "muted",
                "yellow",
            ),
        ),
    ],
}

Ooh, just like HeaderMap in hyper?

Yup! We discussed that recently.

Let’s now try parsing our sample rules into this data structure.

peg?

peg.

$ cargo add peg
      Adding peg v0.6.3 to dependencies

fn parse_rules(input: &str) -> Rules<'_> {
    let mut rules: Rules = Default::default();

    peg::parser! {
        pub(crate) grammar parser() for str {
            pub(crate) rule root(r: &mut Rules<'input>)
                = (line(r) "." whitespace()*)* ![_]

            rule line(r: &mut Rules<'input>)
                = spec:bag_spec() " contain " rules:rules() {
                    if let Some(rules) = rules {
                        for rule in rules {
                            r.insert(spec, rule)
                        }
                    }
                }

            rule bag_spec() -> BagSpec<'input>
                = adjective:name() " " color:name() " bag" "s"? { (adjective, color) }

            rule rules() -> Option<Vec<(usize, BagSpec<'input>)>>
                = rules:rule1()+ { Some(rules) }
                / "no other bags" { None }

            /// Rule followed by an optional comma and space
            rule rule1() -> (usize, BagSpec<'input>)
                = r:rule0() ", "? { r }

            /// A single rule
            rule rule0() -> (usize, BagSpec<'input>)
                = quantity:number() " " spec:bag_spec() { (quantity, spec) }

            rule number() -> usize
                = e:$(['0'..='9']+) { e.parse().unwrap() }

            /// A sequence of non-whitespace characters
            rule name() -> &'input str
                = $((!whitespace()[_])*)

            /// Spaces, tabs, CR and LF
            rule whitespace()
                = [' ' | '\t' | '\r' | '\n']
        }
    }

    parser::root(input, &mut rules).unwrap();
    rules
}

Now, dbg!() makes for pretty verbose output - maybe we can replicate the formatting of the input, for inspection?

use std::fmt;

struct FormattedRules<'a>(Rules<'a>);

impl fmt::Display for FormattedRules<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (k, vv) in &self.0 {
            write!(f, "{} {} bags can contain ", k.0, k.1)?;
            if vv.is_empty() {
                write!(f, "no other bags")?;
            } else {
                for (i, v) in vv.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(
                        f,
                        "{} {} {} {}",
                        v.0,
                        v.1 .0,
                        v.1 .1,
                        if v.0 == 1 { "bag" } else { "bags" }
                    )?;
                }
            }
            writeln!(f, ".")?;
        }
        Ok(())
    }
}

fn main() {
    let rules = parse_rules(include_str!("input.txt"));
    print!("{}", FormattedRules(rules));
}

Now, to compare the output of our program to the sample input, we have to sort it - good old sort command will work out just fine:

$ cargo run --quiet | sort -n
bright white bags can contain 1 shiny gold bag.
dark olive bags can contain 3 faded blue bags, 4 dotted black bags.
dark orange bags can contain 3 bright white bags, 4 muted yellow bags.
light red bags can contain 1 bright white bag, 2 muted yellow bags.
muted yellow bags can contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags can contain 1 dark olive bag, 2 vibrant plum bags.
vibrant plum bags can contain 5 faded blue bags, 6 dotted black bags.

$ cat src/input.txt | sort -n
bright white bags contain 1 shiny gold bag.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
dotted black bags contain no other bags.
faded blue bags contain no other bags.
light red bags contain 1 bright white bag, 2 muted yellow bags.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.

Seems good! The only difference is we don’t remember bags that can contain no other bags. I’m going to say that doesn’t matter right now - but we can always change our mind later.

Back to the question. Even though this data structure is flat, we have a graph on our hands:

It’s not just any graph, either. It’s a Directed Acyclic Graph.

Oh, DAGs! Sure! I like DAGs.

So if we want to know “what bag colors can eventually contain a shiny gold bag?”, we can just walk the graph. For every node in the graph, we just follow the arrows until we either:

Reach “shiny gold” - in which case, yes, it can contain shiny gold bags!
Run out of edges to walk, in which case, no, it cannot contain shiny gold bags.

Let’s take “light red” for example. If we follow the whole graph, these are all the edges and nodes we encounter:

And this subset does contain “shiny gold”, so “light red” bags can eventually contain “shiny gold” bags.

If we start with dark olive however, we only encounter these:

And none of these are “shiny gold”.

Now all we have left is the easy part - make our dreams into reality:

fn subgraph_contains(graph: &Rules<'_>, root: &(&str, &str), needle: &(&str, &str)) -> bool {
    if let Some(neighbors) = graph.get_vec(root) {
        for (_, neighbor) in neighbors {
            if neighbor == needle || subgraph_contains(graph, neighbor, needle) {
                return true;
            }
        }
    }
    false
}

Do we need those if let and that for loop? Couldn’t we just use iterator methods?

Fine, something like this then:

fn subgraph_contains(graph: &Rules<'_>, root: &(&str, &str), needle: &(&str, &str)) -> bool {
    graph
        .get_vec(root)
        .map(|v| {
            v.iter().any(|(_, neighbor)| {
                neighbor == needle || subgraph_contains(graph, neighbor, needle)
            })
        })
        .unwrap_or_default()
}

Ehh couldn’t we flatten this a bit?

Sure, sure, something like that:

fn subgraph_contains(graph: &Rules<'_>, root: &(&str, &str), needle: &(&str, &str)) -> bool {
    graph
        .get_vec(root)
        .unwrap_or(&Default::default())
        .iter()
        .any(|(_, neighbor)| neighbor == needle || subgraph_contains(graph, neighbor, needle))
}

Ah, that’s sorta silly - won’t it allocate in case there there’s no values associated to the root key?

Welp, too late now.

Fine, but do you promise we’ll learn a trick later that could make this better?

Yes. I promise. There!

ANYWAY, we can use such a function like that:

fn main() {
    let rules = parse_rules(include_str!("input.txt"));

    let needle = &("shiny", "gold");
    let colors_that_contain_shiny_gold: Vec<_> = rules
        .keys()
        // shiny gold bags are already shiny god, we're not interested
        // in what they can contain (as per the example)
        .filter(|&k| k != needle)
        .filter(|&k| subgraph_contains(&rules, k, needle))
        .collect();
    println!("{:?}", colors_that_contain_shiny_gold);
}

And we would get that answer:

$ cargo run --quiet
[("dark", "orange"), ("light", "red"), ("bright", "white"), ("muted", "yellow")]

Which matches the example in the problem statement:

In the above rules, the following options would be available to you:

A bright white bag, which can hold your shiny gold bag directly.

A muted yellow bag, which can hold your shiny gold bag directly, plus some other bags.

A dark orange bag, which can hold bright white and muted yellow bags, either of which could then hold your shiny gold bag.

A light red bag, which can hold bright white and muted yellow bags, either of which could then hold your shiny gold bag.

So that’s it? We’re done?

Well…

There’s something else I want to try. When we walk the graph starting from all the nodes, we walk the same subgraphs multiple times.

For example, walking the graph from “light red” and from “dark orange” means visiting “bright white”, “muted yellow”, “faded blue”, and “shiny gold” twice:

We can get rid of all this duplciate work if we just switch it up a little.

Just like gittup, we’re going to make things faster by making the arrows go up.

Imagine our graph looked like this instead:

Now our arrows, or edges, mean “N of this bag can be stored in that bag”.

And if we want to ask the question “which bags can stored shiny gold bags?”, we just need to start with “shiny gold” and walk the entire subgraph.

But how does one go about making the arrows go up?

It’s quite simple really:

fn reverse_graph<'a>(graph: &Rules<'a>) -> Rules<'a> {
    let mut reverse: Rules = Default::default();
    for (&node, neighbors) in graph.iter_all() {
        for &(quantity, neighbor) in neighbors {
            reverse.insert(neighbor, (quantity, node));
        }
    }
    reverse
}

Wait.. can you collect into a MultiMap?

Oh right!

fn reverse_graph<'a>(graph: &Rules<'a>) -> Rules<'a> {
    graph
        .iter_all()
        .map(|(&node, neighbors)| {
            neighbors
                .iter()
                .map(move |&(quantity, neighbor)| (neighbor, (quantity, node)))
        })
        .flatten()
        .collect()
}

Ooh, flatten! Is that new?

Old as the world itself. We need it because .iter_all() returns an Iterator<Item = (tuple, Vec<tuple>)> - because it’s a multimap, remember?

Neat!

And then we make a function to walk a subgraph, right?

We do! But what should it return?

Mhhh an iterator?

Let’s start with a Vec:

fn walk_subgraph<'a>(graph: &Rules<'a>, root: &(&str, &str)) -> Vec<(&'a str, &'a str)> {
    let mut res: Vec<_> = Default::default();
    if let Some(neighbors) = graph.get_vec(root) {
        for &(_quantity, neighbor) in neighbors {
            res.push(neighbor);
            res.extend(walk_subgraph(graph, &neighbor));
        }
    }
    res
}

And see what we get:

fn main() {
    let rules = parse_rules(include_str!("input.txt"));
    let rev_rules = reverse_graph(&rules);

    let colors_that_contain_shiny_gold = walk_subgraph(&rev_rules, &("shiny", "gold"));
    println!("{:?}", colors_that_contain_shiny_gold);
}

$ cargo run --quiet
[("bright", "white"), ("light", "red"), ("dark", "orange"), ("muted", "yellow"), ("light", "red"), ("dark", "orange")]

Hey! There’s duplicates!

Indeed, if we naively walk the graph starting from “shiny gold”, we’re going to visit some nodes more than once.

We could collect into a HashSet, or do something else - for now let’s worry about not allocating a Vec every time we walk a subgraph.

One option is to just take a &mut Vec. That’s completely legitimate:

fn walk_subgraph1<'a>(graph: &Rules<'a>, root: &(&str, &str), res: &mut Vec<(&'a str, &'a str)>) {
    if let Some(neighbors) = graph.get_vec(root) {
        for &(_quantity, neighbor) in neighbors {
            res.push(neighbor);
            walk_subgraph1(graph, &neighbor, res);
        }
    }
}

fn main() {
    let rules = parse_rules(include_str!("input.txt"));
    let rev_rules = reverse_graph(&rules);

    let mut colors_that_contain_shiny_gold = Default::default();
    walk_subgraph1(
        &rev_rules,
        &("shiny", "gold"),
        &mut colors_that_contain_shiny_gold,
    );
    println!("{:?}", colors_that_contain_shiny_gold);
}

This gives us the exact same result, with fewer allocations. Again, it’s a legitimate technique that allows us to not worry about the borrow checker as much.

Of course now we cannot collect into a HashSet to filter out duplicates. And we didn’t even want to do that in the first place, we just needed to count them!

So let’s try making a version that returns an iterator. The problem we’re going to encounter is that our graph can be infinitely large, and so our iterator type, if done naively, will also be infinitely large - I’ve talked about this in May of 2019.

Ah yes, back when your articles still fit inside of a coffee break.

The tl;dr of this article was: Box is your friend, so, without further ado:

fn walk_subgraph2<'iter, 'elems: 'iter>(
    graph: &'iter Rules<'elems>,
    root: &(&'iter str, &'iter str),
) -> Box<dyn Iterator<Item = (&'elems str, &'elems str)> + 'iter> {
    Box::new(
        graph
            .get_vec(root)
            .into_iter()
            .flatten()
            .map(move |&(_, neighbor)| {
                std::iter::once(neighbor).chain(walk_subgraph2(graph, &neighbor))
            })
            .flatten(),
    )
}

Quick question! So the into_iter() I understand - we have a &Vec<V>, which we turn into an Iterator<Item = V>. But what’s with the flatten()?

Well actually… we don’t have a &Vec<T>.

We don’t?

    pub fn get_vec<Q: ?Sized>(&self, k: &Q) -> Option<&Vec<V>>
        where K: Borrow<Q>,
              Q: Eq + Hash

Ohh we have an Option<&Vec<V>>! That’s iterable?

Of course it is! It’s an iterator that yields one element if it’s Some, or zero elements if it’s None.

Wicked. And flatten() gives us an iterator that either yield all the elements of the Vec, or nothing at all?

OHHhhhh is that the trick you mentioned that could’ve replaced .unwrap_or(&Default::default()) with something better?

Yup! 😎

Now that we have an iterator, we can use another goodie from itertools, and this time, I triple-checked that it’s not in the standard library:

$ cargo add itertools
      Adding itertools v0.9.0 to dependencies

use itertools::Itertools;

fn main() {
    let rules = parse_rules(include_str!("input.txt"));
    let rev_rules = reverse_graph(&rules);

    let needle = ("shiny", "gold");
    let answer = walk_subgraph2(&rev_rules, &needle).unique().count();
    println!("{} colors can contain {:?} bags", answer, needle);
}

$ cargo run --quiet
4 colors can contain ("shiny", "gold") bags

Ooh, I saw what you did there - &str implements Debug, so (&str, &str) does too. Neat!

But how does .unique() work?

It… it collects into a HashSet. But that’s hidden! So our code is nice and functional.

Let’s run it on my actual puzzle input:

$ cargo run --quiet
103 colors can contain ("shiny", "gold") bags

Correct!

Part 2

We’re now asked to count how many bags you must buy to fit into a “shiny gold” bag. We’re finally putting those numbers to use!

Time for another walk_subgraph method - one that actually includes quantities.

Only a few changes are needed:

fn walk_subgraph3<'iter, 'elems: 'iter>(
    graph: &'iter Rules<'elems>,
    root: &(&'iter str, &'iter str),
                          // 👇 we're now returning the quantity as well
) -> Box<dyn Iterator<Item = (usize, (&'elems str, &'elems str))> + 'iter> {
    Box::new(
        graph
            .get_vec(root)
            .into_iter()
            .flatten()
            // 👇 this is even simpler, since we're not destructing the tuple anymore
            .map(move |&n| std::iter::once(n).chain(walk_subgraph3(graph, &n.1)))
            .flatten(),
    )
}

We’re going to use the graph with the arrows pointing down, so we don’t need to reverse the rules anymore. Also, we’re not interested in .unique() anymore.

If “shiny gold” bags contain both “dark red” and “light red” bags, and both of these contain “dark teal” bags - we must count “dark teal” bags twice.

So, our main function is even simpler:

fn main() {
    let rules = parse_rules(include_str!("input.txt"));
    let root = ("shiny", "gold");
    let answer = walk_subgraph3(&rules, &root).count();
    println!("you must buy {} bags to fill a {:?} bag", answer, root);
}

$ cargo run --quiet
you must buy 63 bags to fill a ("shiny", "gold") bag

Uh oh… that’s not the right answer!

Don’t worry - it happens to the best of us. Too bad for the streak though.

Oh right - in my excitement, I forgot we need to multiply stuff together!

If every “shiny gold” bag has two “dark red” bags, and those have three “light magenta” bags, then we have 2*3 = 6 “light magenta” bags.

I guess we can’t have a method as generic as walk_subgraph3 - let’s make a custom one.

fn bag_quantities<'iter, 'elems: 'iter>(
    graph: &'iter Rules<'elems>,
    root: &(&'iter str, &'iter str),
) -> Box<dyn Iterator<Item = usize> + 'iter> {
    Box::new(
        graph
            .get_vec(root)
            .into_iter()
            .flatten()
            .map(move |&(qt, n)| {
                std::iter::once(qt).chain(bag_quantities(graph, &n).map(move |x| x * qt))
            })
            .flatten(),
    )
}

fn main() {
    let rules = parse_rules(include_str!("input.txt"));
    let root = ("shiny", "gold");
    let answer: usize = bag_quantities(&rules, &root).sum();
    println!("you must buy {} bags to fill a {:?} bag", answer, root);
}

$ cargo run --quiet
you must buy 1469 bags to fill a ("shiny", "gold") bag

That’s more like it!

I knew 63 was suspiciously low…

Until next time, take care!

You're reading the Advent of Code 2020 series.

Here's another article just for you:

Small strings in Rust

Jul 05, 2020

17 min #rust

Hey everyone!

This article is brought to you by a shameless nerd snipe, courtesy of Pascal.

In case you’ve blocked Twitter for your own good, this reads:

There should be a post explaining and comparing smolstr and smartstring (and maybe others, like smallstr)

Well, I took the bait.

But, since this is me writing, I get to set the rules:

There will be no “maybe others” - we’ll review just the first two

Day 7 (Advent of Code 2020)

1

2

3

4

5

6

7

8

9

10

11

12

13

14

Part 2

1

2

3

4

5

6

7

8

9

10

11

12

13

14

Small strings in Rust