use crate::input::Input;
pub fn solve(input: &Input) -> Result<i64, String> {
let iterations = input.part_values(1, 10);
let input_multiplier = input.part_values(1, 811_589_153);
let numbers = input
.text
.lines()
.map(|line| Some(i64::from(line.parse::<i16>().ok()?) * input_multiplier))
.collect::<Option<Vec<_>>>()
.ok_or("Invalid input - lines are not numbers in the range [-32,768, 32_767]")?;
if numbers.len() < 3 {
return Err("Input must have at least three numbers".to_string());
} else if numbers.len() > MAX_LENGTH {
return Err(format!("Too many numbers - max {MAX_LENGTH}"));
}
let zero_idx = numbers
.iter()
.position(|&n| n == 0)
.ok_or("No zero value in input")?;
let bucket_size = (numbers.len() as f64).sqrt() as usize;
let mut buckets = Vec::with_capacity(numbers.len() / bucket_size);
for i in (0..numbers.len()).step_by(bucket_size) {
let range_end = (i + bucket_size).min(numbers.len());
buckets.push((i..range_end).collect::<Vec<_>>());
}
let mut idx_to_bucket = (0..numbers.len())
.map(|i| i / bucket_size)
.collect::<Vec<_>>();
for _ in 0..iterations {
for idx_to_shift in 0..numbers.len() {
let shift_at_idx =
(numbers[idx_to_shift]).rem_euclid(numbers.len() as i64 - 1) as usize;
let old_bucket_containing_idx = idx_to_bucket[idx_to_shift];
let old_offset_in_bucket = buckets[old_bucket_containing_idx]
.iter()
.position(|&n| n == idx_to_shift)
.unwrap_or_default();
buckets[old_bucket_containing_idx].remove(old_offset_in_bucket);
let (new_bucket_containing_idx, new_offset_in_bucket) = find_bucket_and_offset(
&buckets,
old_bucket_containing_idx,
old_offset_in_bucket + shift_at_idx,
numbers.len(),
);
buckets[new_bucket_containing_idx].insert(new_offset_in_bucket, idx_to_shift);
idx_to_bucket[idx_to_shift] = new_bucket_containing_idx;
}
}
let mut bucket_containing_number = idx_to_bucket[zero_idx];
let mut number_offset_in_bucket = buckets[bucket_containing_number]
.iter()
.position(|&n| n == zero_idx)
.unwrap_or_default();
Ok(std::iter::from_fn(|| {
(bucket_containing_number, number_offset_in_bucket) = find_bucket_and_offset(
&buckets,
bucket_containing_number,
number_offset_in_bucket + 1000,
numbers.len(),
);
let number_idx = buckets[bucket_containing_number][number_offset_in_bucket];
Some(numbers[number_idx])
})
.take(3)
.sum())
}
fn find_bucket_and_offset(
buckets: &[Vec<usize>],
mut bucket: usize,
mut offset_relative_to_bucket: usize,
len: usize,
) -> (usize, usize) {
if offset_relative_to_bucket < len / 2
|| (offset_relative_to_bucket > (len - 1 + buckets[bucket].len()))
{
while offset_relative_to_bucket >= buckets[bucket].len() {
offset_relative_to_bucket -= buckets[bucket].len();
bucket = (bucket + 1) % buckets.len();
}
(bucket, offset_relative_to_bucket)
} else {
offset_relative_to_bucket = len + buckets[bucket].len() - 1 - offset_relative_to_bucket;
while offset_relative_to_bucket >= buckets[bucket].len() {
offset_relative_to_bucket -= buckets[bucket].len();
bucket = bucket.checked_sub(1).unwrap_or(buckets.len() - 1);
}
(bucket, buckets[bucket].len() - offset_relative_to_bucket)
}
}
const MAX_LENGTH: usize = 10_000;
#[test]
pub fn tests() {
use crate::input::{test_part_one, test_part_one_error, test_part_two};
let test_input = "1\n2\n-3\n3\n-2\n0\n4";
test_part_one!(test_input => 3);
test_part_two!(test_input => 1_623_178_306);
let real_input = include_str!("day20_input.txt");
test_part_one!(real_input => 7_225);
test_part_two!(real_input => 548_634_267_428);
let test_input = "0\n1\n2";
test_part_one!(test_input => 3);
let test_input = "0\n1";
test_part_one_error!(test_input => "Input must have at least three numbers");
}