|
| 1 | +use aoc_utils::*; |
| 2 | +use itertools::Itertools; |
| 3 | +use z3::{ |
| 4 | + ast::{Ast, Int}, |
| 5 | + Config, Context, Optimize, SatResult, |
| 6 | +}; |
| 7 | + |
| 8 | +advent_of_code::solution!(10); |
| 9 | + |
| 10 | +fn parse_lights(input: &str) -> u64 { |
| 11 | + let mut lights = 0; |
| 12 | + let s = input.trim_start_matches('[').trim_end_matches(']'); |
| 13 | + for (i, c) in s.as_bytes().iter().enumerate() { |
| 14 | + if *c == b'#' { |
| 15 | + lights |= 1 << i; |
| 16 | + } |
| 17 | + } |
| 18 | + lights |
| 19 | +} |
| 20 | + |
| 21 | +fn parse_button(input: &str) -> u64 { |
| 22 | + let mut xor = 0; |
| 23 | + let s = input |
| 24 | + .trim_start_matches('(') |
| 25 | + .trim_end_matches(')') |
| 26 | + .split(','); |
| 27 | + for c in s.map(|c| c.as_bytes()) { |
| 28 | + assert!(c.len() == 1); |
| 29 | + let d = c[0] - b'0'; |
| 30 | + xor |= 1 << d; |
| 31 | + } |
| 32 | + xor |
| 33 | +} |
| 34 | + |
| 35 | +fn solve_part1(target: u64, buttons: &[u64]) -> Option<u64> { |
| 36 | + for count in 1..=buttons.len() { |
| 37 | + for v in buttons.iter().permutations(count) { |
| 38 | + let mut curr = 0; |
| 39 | + v.into_iter().for_each(|x| curr ^= x); |
| 40 | + if curr == target { |
| 41 | + return Some(count as u64); |
| 42 | + } |
| 43 | + } |
| 44 | + } |
| 45 | + None |
| 46 | +} |
| 47 | + |
| 48 | +pub fn part_one(input: &str) -> Option<u64> { |
| 49 | + // gonna use bits for everything because xor can work here |
| 50 | + let lines = input.mlines(|s| { |
| 51 | + let mut i = s.split_whitespace(); |
| 52 | + let lights = i.next().map(parse_lights).expect("Didn't find lights"); |
| 53 | + // don't need joltages for part 1 |
| 54 | + let _ = i.next_back().expect("Didn't find joltages"); |
| 55 | + let buttons = i.map(parse_button).collect_vec(); |
| 56 | + (lights, buttons) |
| 57 | + }); |
| 58 | + Some( |
| 59 | + lines |
| 60 | + .iter() |
| 61 | + .map(|(target, buttons)| { |
| 62 | + solve_part1(*target, buttons).expect("Should find solution for machine") |
| 63 | + }) |
| 64 | + .sum(), |
| 65 | + ) |
| 66 | +} |
| 67 | + |
| 68 | +fn solve_part2(buttons: &[Vec<usize>], joltages: &[u64]) -> u64 { |
| 69 | + let cfg = Config::new(); |
| 70 | + let ctx = Context::new(&cfg); |
| 71 | + let opt = Optimize::new(&ctx); |
| 72 | + |
| 73 | + // target_counts = number of times to press each button (what we're solving for) |
| 74 | + let target_counts: Vec<Int> = (0..buttons.len()) |
| 75 | + .map(|i| Int::new_const(&ctx, format!("p{i}"))) |
| 76 | + .collect(); |
| 77 | + |
| 78 | + // Constraint: (sum of target_counts[i] for all i where light_pos in buttons[i]) == joltages[light_pos] |
| 79 | + for (light_pos, &jolt) in joltages.iter().enumerate() { |
| 80 | + let terms: Vec<_> = buttons |
| 81 | + .iter() |
| 82 | + .enumerate() |
| 83 | + .filter(|(_, b)| b.contains(&light_pos)) |
| 84 | + .map(|(i, _)| &target_counts[i]) |
| 85 | + .collect(); |
| 86 | + |
| 87 | + let sum = Int::add(&ctx, &terms); |
| 88 | + opt.assert(&sum._eq(&Int::from_u64(&ctx, jolt))); |
| 89 | + } |
| 90 | + |
| 91 | + // Constraint: t >= 0 (can't press a button negative times) |
| 92 | + let zero = Int::from_u64(&ctx, 0); |
| 93 | + for t in &target_counts { |
| 94 | + opt.assert(&t.ge(&zero)); |
| 95 | + } |
| 96 | + |
| 97 | + // Minimize sum of target_counts |
| 98 | + let target_count_refs: Vec<_> = target_counts.iter().collect(); |
| 99 | + let total = Int::add(&ctx, &target_count_refs); |
| 100 | + opt.minimize(&total); |
| 101 | + |
| 102 | + // Solve and extract result |
| 103 | + assert_eq!(opt.check(&[]), SatResult::Sat); |
| 104 | + let model = opt.get_model().unwrap(); |
| 105 | + |
| 106 | + target_counts |
| 107 | + .iter() |
| 108 | + .map(|v| model.eval(v, true).unwrap().as_u64().unwrap()) |
| 109 | + .sum() |
| 110 | +} |
| 111 | + |
| 112 | +fn parse_buttons_part2(input: &str) -> Vec<usize> { |
| 113 | + // remove () before split |
| 114 | + input |
| 115 | + .trim_start_matches('(') |
| 116 | + .trim_end_matches(')') |
| 117 | + .split(',') |
| 118 | + .map(|x| x.parse().expect("Button item should be int")) |
| 119 | + .collect_vec() |
| 120 | +} |
| 121 | + |
| 122 | +fn parse_joltages(input: &str) -> Vec<u64> { |
| 123 | + // remove {} before split |
| 124 | + input |
| 125 | + .trim_start_matches('{') |
| 126 | + .trim_end_matches('}') |
| 127 | + .split(',') |
| 128 | + .map(|x| x.parse().expect("Button item should be int")) |
| 129 | + .collect_vec() |
| 130 | +} |
| 131 | + |
| 132 | +pub fn part_two(input: &str) -> Option<u64> { |
| 133 | + // gonna use bits for everything because xor can work here |
| 134 | + let lines = input.mlines(|s| { |
| 135 | + let mut i = s.split_whitespace(); |
| 136 | + // don't need lights for part 2 |
| 137 | + let _ = i.next().expect("Didn't find lights"); |
| 138 | + let joltages = i |
| 139 | + .next_back() |
| 140 | + .map(parse_joltages) |
| 141 | + .expect("Didn't find joltages"); |
| 142 | + let buttons = i.map(parse_buttons_part2).collect_vec(); |
| 143 | + (buttons, joltages) |
| 144 | + }); |
| 145 | + Some(lines.iter().map(|(b, j)| solve_part2(b, j)).sum()) |
| 146 | +} |
| 147 | + |
| 148 | +#[cfg(test)] |
| 149 | +mod tests { |
| 150 | + use super::*; |
| 151 | + |
| 152 | + #[test] |
| 153 | + fn test_part_one() { |
| 154 | + let result = part_one(&advent_of_code::template::read_file("examples", DAY)); |
| 155 | + assert_eq!(result, Some(7)); |
| 156 | + } |
| 157 | + |
| 158 | + #[test] |
| 159 | + fn test_part_two() { |
| 160 | + let result = part_two(&advent_of_code::template::read_file("examples", DAY)); |
| 161 | + assert_eq!(result, Some(33)); |
| 162 | + } |
| 163 | +} |
0 commit comments