AdventOfCode/Rust/2023/24.rs

151 lines
4.6 KiB
Rust

#![feature(test)]
use std::ops::{Add, Mul};
use itertools::{Itertools, PeekingNext};
use num::rational::Ratio;
use rayon::prelude::*;
use z3::{
ast::{Ast, Int, Real},
Config, Context, Solver,
};
#[derive(Debug)]
struct Input {
hailstones: Vec<Hailstone>,
test_start: i64,
test_end: i64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Hailstone {
position: Vec3,
velocity: Vec3,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Vec3 {
x: i64,
y: i64,
z: i64,
}
fn setup(input: &str) -> Input {
let mut lines = input.lines().peekable();
let (test_start, test_end) = lines
.peeking_next(|line| line.starts_with('#'))
.map(|line| {
line.split(' ')
.skip(1)
.map(|x| x.parse().unwrap())
.collect_tuple()
.unwrap()
})
.unwrap_or((200000000000000, 400000000000000));
let hailstones = lines
.map(|line| {
let (position, velocity) = line
.split('@')
.map(|split| {
let (x, y, z) = split
.split(',')
.map(|x| x.trim().parse().unwrap())
.collect_tuple()
.unwrap();
Vec3 { x, y, z }
})
.collect_tuple()
.unwrap();
Hailstone { position, velocity }
})
.collect();
Input {
hailstones,
test_start,
test_end,
}
}
fn part1(input: &Input) -> usize {
let test_start = Ratio::from(input.test_start as i128);
let test_end = Ratio::from(input.test_end as i128);
input
.hailstones
.iter()
.copied()
.enumerate()
.flat_map(|(i, a)| {
input
.hailstones
.iter()
.copied()
.take(i)
.map(move |b| (a, b))
})
.filter(|&(a, b)| a.velocity.x * b.velocity.y != a.velocity.y * b.velocity.x)
.par_bridge()
.filter(|&(a, b)| {
let x1 = Ratio::from(a.position.y as i128);
let y1 = Ratio::from(a.position.x as i128);
let x2 = Ratio::from(b.position.y as i128);
let y2 = Ratio::from(b.position.x as i128);
let vx1 = Ratio::from(a.velocity.y as i128);
let vy1 = Ratio::from(a.velocity.x as i128);
let vx2 = Ratio::from(b.velocity.y as i128);
let vy2 = Ratio::from(b.velocity.x as i128);
let t1 = ((y1 - y2) * vx2 - (x1 - x2) * vy2) / (vx1 * vy2 - vy1 * vx2);
let t2 = ((y2 - y1) * vx1 - (x2 - x1) * vy1) / (vx2 * vy1 - vy2 * vx1);
let x = x1 + t1 * vx1;
let y = y1 + t1 * vy1;
debug_assert_eq!(x, x2 + t2 * vx2);
debug_assert_eq!(y, y2 + t2 * vy2);
t1 >= 0.into()
&& t2 >= 0.into()
&& x >= test_start
&& x <= test_end
&& y >= test_start
&& y <= test_end
})
.count()
}
fn part2(input: &Input) -> i64 {
let config = Config::new();
let ctx = Context::new(&config);
let solver = Solver::new(&ctx);
let x = Real::new_const(&ctx, "x");
let y = Real::new_const(&ctx, "y");
let z = Real::new_const(&ctx, "z");
let vx = Real::new_const(&ctx, "vx");
let vy = Real::new_const(&ctx, "vy");
let vz = Real::new_const(&ctx, "vz");
let sum = Real::new_const(&ctx, "sum");
solver.assert(&(&x).add(&y).add(&z)._eq(&sum));
for (i, hailstone) in input.hailstones.iter().take(3).copied().enumerate() {
let ti = Real::new_const(&ctx, format!("t{i}"));
let xi = Real::from_int(&Int::from_i64(&ctx, hailstone.position.x));
let yi = Real::from_int(&Int::from_i64(&ctx, hailstone.position.y));
let zi = Real::from_int(&Int::from_i64(&ctx, hailstone.position.z));
let vxi = Real::from_int(&Int::from_i64(&ctx, hailstone.velocity.x));
let vyi = Real::from_int(&Int::from_i64(&ctx, hailstone.velocity.y));
let vzi = Real::from_int(&Int::from_i64(&ctx, hailstone.velocity.z));
solver.assert(&vxi.mul(&ti).add(&xi)._eq(&(&vx).mul(&ti).add(&x)));
solver.assert(&vyi.mul(&ti).add(&yi)._eq(&(&vy).mul(&ti).add(&y)));
solver.assert(&vzi.mul(&ti).add(&zi)._eq(&(&vz).mul(&ti).add(&z)));
}
solver.check();
let model = solver.get_model().unwrap();
let sum = model.get_const_interp(&sum).unwrap();
sum.as_real().unwrap().0
}
aoc::main!(2023, 24, ex: 1);