2020: day16: part 2

aoc2020/input/day16_provided2.txt

@@ -0,0 +1,11 @@
+class: 0-1 or 4-19
+row: 0-5 or 8-19
+seat: 0-13 or 16-19
+
+your ticket:
+11,12,13
+
+nearby tickets:
+3,9,18
+15,1,5
+5,14,9

aoc2020/src/day16.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::fmt::Write;
 use std::ops::RangeInclusive;
 
@@ -9,6 +10,7 @@ pub fn run() -> Result<String> {
     let mut res = String::with_capacity(128);
 
     writeln!(res, "part 1: {}", part1(INPUT)?)?;
+    writeln!(res, "part 2: {}", part2(INPUT)?)?;
 
     Ok(res)
 }
@@ -16,14 +18,88 @@ pub fn run() -> Result<String> {
 fn part1(input: &str) -> Result<u64> {
     let (fields, _, tickets) = parse_input(input)?;
 
+    let fields_vec = fields.values().collect::<Vec<_>>();
+
     Ok(tickets
         .iter()
-        .map(|t| t.invalid_values(&fields))
+        .map(|t| t.invalid_values(&fields_vec))
         .flatten()
         .sum())
 }
 
-fn parse_input(input: &str) -> Result<(Vec<Field>, Ticket, Vec<Ticket>)> {
+fn assign_field_positions(
+    fields: HashMap<&str, Field>,
+    tickets: Vec<Ticket>,
+) -> HashMap<usize, &str> {
+    let fields_vec = fields.values().collect::<Vec<_>>();
+
+    let tickets: Vec<Ticket> = tickets
+        .into_iter()
+        .filter(|t| t.invalid_values(&fields_vec).count() == 0)
+        .collect();
+
+    let num_values = tickets[0].values.len();
+
+    // build list of all possibilities for each field
+    let mut possibilities = fields
+        .iter()
+        .map(|(&name, field)| {
+            let possibilities = (0..num_values)
+                .into_iter()
+                .filter(|i| tickets.iter().all(|t| t.valid_field(field, *i)))
+                .collect();
+
+            (name, possibilities)
+        })
+        .collect::<HashMap<&str, Vec<usize>>>();
+
+    let mut fields_to_assign: Vec<&str> = fields.keys().copied().collect();
+    let mut field_indices: HashMap<usize, &str> = HashMap::new();
+
+    for _ in 0..fields.len() {
+        // get field with only one possibility, and assign it
+        let field = fields_to_assign
+            .iter()
+            .min_by_key(|&name| possibilities[name].len())
+            .copied()
+            .expect("fields_to_assign should never be empty in this loop");
+
+        assert_eq!(possibilities[field].len(), 1);
+
+        let possibility = possibilities[field][0];
+
+        // remove position from other fields' possibilities
+        possibilities.values_mut().for_each(|list| {
+            if let Some(pos) = list.iter().position(|idx| *idx == possibility) {
+                list.swap_remove(pos);
+            }
+        });
+
+        field_indices.insert(possibility, field);
+        fields_to_assign.swap_remove(fields_to_assign.iter().position(|f| *f == field).unwrap());
+    }
+
+    field_indices
+}
+
+fn part2(input: &str) -> Result<u64> {
+    let (fields, my_ticket, tickets) = parse_input(input)?;
+
+    let field_pos_matches = assign_field_positions(fields, tickets);
+
+    Ok(my_ticket
+        .values
+        .iter()
+        .enumerate()
+        .filter(|(idx, _)| {
+            let field = field_pos_matches[idx];
+            field.starts_with("departure")
+        })
+        .map(|(_, val)| val)
+        .product())
+}
+
+fn parse_input(input: &str) -> Result<(HashMap<&str, Field>, Ticket, Vec<Ticket>)> {
     let mut parts = input.split("\n\n");
 
     let fields_part = parts.next().context("no fields specification found")?;
@@ -32,7 +108,13 @@ fn parse_input(input: &str) -> Result<(Vec<Field>, Ticket, Vec<Ticket>)> {
 
     let fields = fields_part
         .lines()
-        .map(|line| line.parse())
+        .map(|line| {
+            let mut parts = line.split(": ");
+            let name = parts.next().context("no name found")?;
+            let field = parts.next().context("no ranges found")?.parse()?;
+
+            Ok((name, field))
+        })
         .collect::<Result<_>>()
         .context("couldn't parse fields")?;
     let my_ticket = my_ticket_part
@@ -57,12 +139,15 @@ struct Ticket {
 }
 
 impl Ticket {
-    fn invalid_values(&self, fields: &[Field]) -> Vec<u64> {
+    fn invalid_values<'a>(&'a self, fields: &'a [&Field]) -> impl Iterator<Item = u64> + 'a {
         self.values
             .iter()
-            .filter(|val| !fields.iter().any(|f| f.contains(*val)))
             .copied()
-            .collect()
+            .filter(move |val| !fields.iter().any(|f| f.contains(val)))
+    }
+
+    fn valid_field(&self, field: &Field, val_idx: usize) -> bool {
+        field.contains(&self.values[val_idx])
     }
 }
 
@@ -84,7 +169,6 @@ impl std::str::FromStr for Ticket {
 
 #[derive(Debug)]
 struct Field {
-    name: String,
     ranges: (RangeInclusive<u64>, RangeInclusive<u64>),
 }
 
@@ -98,11 +182,7 @@ impl std::str::FromStr for Field {
     type Err = anyhow::Error;
 
     fn from_str(s: &str) -> Result<Self> {
-        let mut parts = s.split(": ");
-
-        let name = parts.next().context("no name found")?.to_string();
-
-        let mut ranges = parts.next().context("no ranges found")?.split(" or ");
+        let mut ranges = s.split(" or ");
 
         let mut range1 = ranges.next().context("no first range found")?.split('-');
         let range1_start = range1.next().context("no bound for range")?.parse()?;
@@ -113,7 +193,6 @@ impl std::str::FromStr for Field {
         let range2_end = range2.next().context("no bound for range")?.parse()?;
 
         Ok(Field {
-            name,
             ranges: (range1_start..=range1_end, range2_start..=range2_end),
         })
     }
@@ -123,15 +202,36 @@ impl std::str::FromStr for Field {
 mod tests {
     use super::*;
 
-    const PROVIDED: &str = include_str!("../input/day16_provided.txt");
+    const PROVIDED1: &str = include_str!("../input/day16_provided1.txt");
+    const PROVIDED2: &str = include_str!("../input/day16_provided2.txt");
 
     #[test]
     fn part1_provided() {
-        assert_eq!(part1(PROVIDED).unwrap(), 71);
+        assert_eq!(part1(PROVIDED1).unwrap(), 71);
     }
 
     #[test]
     fn part1_real() {
         assert_eq!(part1(INPUT).unwrap(), 20013);
     }
+
+    #[test]
+    fn part2_provided() {
+        let (fields, _, tickets) = parse_input(PROVIDED2).unwrap();
+
+        let matches = assign_field_positions(fields, tickets);
+
+        let expected = (&["row", "class", "seat"])
+            .iter()
+            .copied()
+            .enumerate()
+            .collect();
+
+        assert_eq!(matches, expected);
+    }
+
+    #[test]
+    fn part2_real() {
+        assert_eq!(part2(INPUT).unwrap(), 5977293343129);
+    }
 }