aoc2019/day12/src/main.rs

fn strip_prefix<'a>(prefix: &str, input: &'a str) -> &'a str {
	assert!(input.starts_with(prefix));
	&input[prefix.len()..]
}

fn strip_suffix<'a>(suffix: &str, input: &'a str) -> &'a str {
	assert!(input.ends_with(suffix));
	&input[..input.len() - suffix.len()]
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Default)]
pub struct Position {
	pub x: i32,
	pub y: i32,
	pub z: i32,
}

impl Position {
	pub fn parse(input: &str) -> Self {
		let input = strip_suffix(">", strip_prefix("<", input.trim()));
		let mut parts = input.split(",");
		let x = strip_prefix("x=", parts.next().expect("need x component").trim()).parse().unwrap();
		let y = strip_prefix("y=", parts.next().expect("need y component").trim()).parse().unwrap();
		let z = strip_prefix("z=", parts.next().expect("need z component").trim()).parse().unwrap();
		assert!(parts.next().is_none(), "unexpected data");
		Position { x, y, z }
	}

	pub fn potential_energy(self) -> u32 {
		self.x.abs() as u32 + self.y.abs() as u32 + self.z.abs() as u32
	}
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Default)]
pub struct Velocity {
	pub dx: i32,
	pub dy: i32,
	pub dz: i32,
}

impl Velocity {
	pub fn kinetic_energy(self) -> u32 {
		self.dx.abs() as u32 + self.dy.abs() as u32 + self.dz.abs() as u32
	}
}

impl std::ops::AddAssign<Velocity> for Position {
	fn add_assign(&mut self, rhs: Velocity) {
		self.x += rhs.dx;
		self.y += rhs.dy;
		self.z += rhs.dz;
	}
}

fn parse_moons(input: &str) -> Vec<(Position, Velocity)> {
	input.lines().map(Position::parse).map(|p| (p, Velocity::default())).collect()
}

fn step(moons: &mut [(Position, Velocity)]) {
	for i in 0..moons.len() - 1 {
		// change velocity based on "gravity"
		for j in i + 1..moons.len() {
			let dx = (moons[j].0.x - moons[i].0.x).signum();
			moons[i].1.dx += dx; moons[j].1.dx -= dx;
			let dy = (moons[j].0.y - moons[i].0.y).signum();
			moons[i].1.dy += dy; moons[j].1.dy -= dy;
			let dz = (moons[j].0.z - moons[i].0.z).signum();
			moons[i].1.dz += dz; moons[j].1.dz -= dz;
		}
	}
	for moon in moons {
		moon.0 += moon.1;
	}
}

fn energy(moons: &[(Position, Velocity)]) -> u32 {
	moons.iter().map(|(p, v)| p.potential_energy() * v.kinetic_energy()).sum()
}

fn energy_after_steps(mut moons: Vec<(Position, Velocity)>, steps: u32) -> u32 {
	for _ in 0..steps {
		step(&mut moons);
	}
	energy(&moons)
}

type ExtractedAxis = Vec<(i32, i32)>;
fn extract_x(moons: &[(Position, Velocity)]) -> ExtractedAxis {
	moons.iter().map(|(p, v)| (p.x, v.dx)).collect()
}
fn extract_y(moons: &[(Position, Velocity)]) -> ExtractedAxis {
	moons.iter().map(|(p, v)| (p.y, v.dy)).collect()
}
fn extract_z(moons: &[(Position, Velocity)]) -> ExtractedAxis {
	moons.iter().map(|(p, v)| (p.z, v.dz)).collect()
}

fn repeats_after_steps(mut moons: Vec<(Position, Velocity)>) -> u64 {
	// as `step` is reversible the first repeated state will always be the initial state

	// find cycle per axis
	let mut result_x = None;
	let initial_x = extract_x(&moons);
	let mut result_y = None;
	let initial_y = extract_y(&moons);
	let mut result_z = None;
	let initial_z = extract_z(&moons);

	for current in 1u64.. {
		step(&mut moons);
		if result_x.is_none() && initial_x == extract_x(&moons) {
			result_x = Some(current);
		}
		if result_y.is_none() && initial_y == extract_y(&moons) {
			result_y = Some(current);
		}
		if result_z.is_none() && initial_z == extract_z(&moons) {
			result_z = Some(current);
		}
		if let (Some(x), Some(y), Some(z)) = (result_x, result_y, result_z) {
			return num_integer::lcm(num_integer::lcm(x, y), z);
		}
	}
	unreachable!()
}

fn main() {
	let moons = parse_moons(include_str!("input.txt"));
	let energy = energy_after_steps(moons.clone(), 1000);
	println!("Energy after 1000 steps: {}", energy);
	println!("Repeats initial state after {} steps", repeats_after_steps(moons));
}

#[cfg(test)]
mod day12test {
	use super::*;

	#[test]
	fn examples1() {
		assert_eq!(179, energy_after_steps(parse_moons(
			"<x=-1, y=0, z=2>
<x=2, y=-10, z=-7>
<x=4, y=-8, z=8>
<x=3, y=5, z=-1>"), 10));
		assert_eq!(1940, energy_after_steps(parse_moons(
			"<x=-8, y=-10, z=0>
<x=5, y=5, z=10>
<x=2, y=-7, z=3>
<x=9, y=-8, z=-3>
"), 100));
	}

	#[test]
	fn examples2() {
		assert_eq!(2772, repeats_after_steps(parse_moons(
			"<x=-1, y=0, z=2>
<x=2, y=-10, z=-7>
<x=4, y=-8, z=8>
<x=3, y=5, z=-1>")));
		assert_eq!(4686774924, repeats_after_steps(parse_moons(
			"<x=-8, y=-10, z=0>
<x=5, y=5, z=10>
<x=2, y=-7, z=3>
<x=9, y=-8, z=-3>")));
	}
}