AoC2020/src/day10.rs

100 lines
2.8 KiB
Rust

use std::io::BufRead;
use std::collections::HashMap;
fn prepare_adapters(adapters: &mut Vec<u32>) {
adapters.sort();
let device = adapters.last().unwrap() + 3;
adapters.push(device);
adapters.insert(0, 0);
}
fn compute_jortage(mut adapters: Vec<u32>) -> (u32, u32) {
prepare_adapters(&mut adapters);
adapters.windows(2)
.map(|window| window[1] - window[0])
.fold((0, 0), |(one, three), diff|
match diff {
1 => (one + 1, three),
3 => (one, three + 1),
_ => (one, three)
})
}
fn count_arrangement(mut adapters: Vec<u32>) -> u64 {
prepare_adapters(&mut adapters);
count_arrangement_rec(adapters.as_slice(), 0, &mut HashMap::new())
}
fn count_arrangement_rec(adapters: &[u32], idx: usize, cache: &mut HashMap<usize, u64>) -> u64 {
if adapters.len() == 1 {
1
} else {
if cache.contains_key(&idx) {
cache[&idx]
} else {
let max_jort = adapters[0] + 3;
let max_idx = if adapters.len() < 4 { adapters.len() } else { 4 };
let count = (1..max_idx).filter_map(|i| {
if adapters[i] > max_jort {
None
} else {
Some(count_arrangement_rec(&adapters[i..], idx + i, cache))
}
}).sum();
cache.insert(idx, count);
count
}
}
}
pub fn part1<F: BufRead> (input: F) {
let input: Vec<_> = input.lines()
.filter_map(|line| line.unwrap().parse::<u32>().ok())
.collect();
let (one, three) = compute_jortage(input);
println!("{}", one * three);
}
pub fn part2<F: BufRead> (input: F) {
let input: Vec<_> = input.lines()
.filter_map(|line| line.unwrap().parse::<u32>().ok())
.collect();
println!("{}", count_arrangement(input));
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_compute_jortage1() {
let input = vec![16, 10, 15, 5, 1, 11, 7, 19, 6, 12, 4];
assert_eq!(compute_jortage(input), (7, 5));
}
#[test]
fn test_compute_jortage2() {
let input = vec![
28, 33, 18, 42, 31, 14, 46, 20, 48, 47, 24, 23, 49, 45, 19,
38, 39, 11, 1, 32, 25, 35, 8, 17, 7, 9, 4, 2, 34, 10, 3
];
assert_eq!(compute_jortage(input), (22, 10));
}
#[test]
fn test_count_arrangement1() {
let input = vec![16, 10, 15, 5, 1, 11, 7, 19, 6, 12, 4];
assert_eq!(count_arrangement(input), 8);
}
#[test]
fn test_count_arrangement2() {
let input = vec![
28, 33, 18, 42, 31, 14, 46, 20, 48, 47, 24, 23, 49, 45, 19,
38, 39, 11, 1, 32, 25, 35, 8, 17, 7, 9, 4, 2, 34, 10, 3
];
assert_eq!(count_arrangement(input), 19208);
}
}