day23.rs

use aoc_runner_derive::aoc;
use aoc_runner_derive::aoc_generator;
use regex::Regex;
use std::cmp;
use rand::prelude::*;

#[derive(Debug, Clone)]
pub struct NanoBot {
    x: i64,
    y: i64,
    z: i64,
    r: i64,
}


#[aoc_generator(day23)]
pub fn input_generator(input: &str) -> Vec<NanoBot> {
    let re = Regex::new(r"pos=<(-?\d+),(-?\d+),(-?\d+)>, r=(\d+)").unwrap();

    input.lines()
        .map(|line| {
            let c = re.captures(line).unwrap();

            NanoBot {
                x: c.get(1).unwrap().as_str().parse().unwrap(),
                y: c.get(2).unwrap().as_str().parse().unwrap(),
                z: c.get(3).unwrap().as_str().parse().unwrap(),
                r: c.get(4).unwrap().as_str().parse().unwrap(),
            }
        })
        .collect()
}

fn dist(lhs: &NanoBot, rhs: &NanoBot) -> i64 {
    (lhs.x - rhs.x).abs() + (lhs.y - rhs.y).abs() + (lhs.z - rhs.z).abs()
}

#[aoc(day23, part1)]
pub fn solve_part1(input: &Vec<NanoBot>) -> usize {
    //println!("{:?}", input);

    let b = input.iter().max_by(|l, r| l.r.cmp(&r.r)).unwrap();

    println!("{:?}", b);

    let in_range = input.iter()
        .filter(|x| dist(x, &b) <= b.r)
        .count();

    in_range
}

pub fn find_most_overlapping(input: &Vec<NanoBot>, in_range: &Vec<Vec<bool>>, remaining_input: &[usize], working_set: &mut Vec<usize>, best: &mut Vec<usize>) {
    for i in 0..remaining_input.len() {
        if best.len() > remaining_input.len()-i+working_set.len() {
            break;
        }

        // check if it fits in
        if !working_set.iter().all(|&r| in_range[r][remaining_input[i]]) {
            continue;
        }

        working_set.push(remaining_input[i].clone());
        find_most_overlapping(input, in_range, &remaining_input[i+1..], working_set, best);
        working_set.pop();
    }

    if working_set.len() > best.len() {
        let loc = find_location(input, working_set);
        println!("{} : {:?}", working_set.len(), loc);
        if let Some(loc) = loc {
            println!("{}", loc.x.abs() + loc.y.abs() + loc.z.abs());
        }
        *best = working_set.clone();
    }
}

fn gen_range(rng: &mut rand::prelude::ThreadRng, max: i64) -> i64 {
    if max > 1 {
        rng.gen_range(1, max)
    } else {
        1
    }
}

fn find_location(input: &Vec<NanoBot>, best: &Vec<usize>) -> Option<NanoBot> {
    let mut loc = input[best[0]].clone();
    loc.r = 0;
    loop {
        let mut dist_moved = 0;
        for out_of_range in best.iter() {
            let target = &input[*out_of_range];

            if dist(&loc, &target) <= target.r {
                continue;
            }

            let mut dist_to_move = dist(&loc, &target) - target.r;
            dist_moved = cmp::max(dist_moved, dist_to_move);
            while dist_to_move > 0 {
                let max_x = (target.x - loc.x).abs();
                let max_y = (target.y - loc.y).abs();
                let max_z = (target.z - loc.z).abs();
                let num_nonzero = (max_x > 0) as i64 + (max_y > 0) as i64 + (max_z > 0) as i64;
                //println!("need to move {} from {:?} to {:?} ", dist_to_move, loc, target);

                if max_x > 0 && (max_x <= max_y || max_y == 0) && (max_x <= max_z || max_z == 0) {
                    let move_amt = cmp::min(max_x, (dist_to_move + num_nonzero - 1) / num_nonzero);
                    loc.x += (target.x - loc.x).signum() * move_amt;
                    loc.y += (target.y - loc.y).signum() * move_amt;
                    loc.z += (target.z - loc.z).signum() * move_amt;
                } else if max_y > 0 && (max_y <= max_z || max_z == 0) {
                    let move_amt = cmp::min(max_y, (dist_to_move + num_nonzero - 1) / num_nonzero);
                    loc.x += (target.x - loc.x).signum() * move_amt;
                    loc.y += (target.y - loc.y).signum() * move_amt;
                    loc.z += (target.z - loc.z).signum() * move_amt;
                } else if max_z > 0 {
                    let move_amt = cmp::min(max_z, (dist_to_move + num_nonzero - 1) / num_nonzero);
                    loc.x += (target.x - loc.x).signum() * move_amt;
                    loc.y += (target.y - loc.y).signum() * move_amt;
                    loc.z += (target.z - loc.z).signum() * move_amt;
                } else {
                    panic!("not all three should be zero");
                }
                dist_to_move = dist(&loc, &target) - target.r;
            }
        }

        if dist_moved < 3 {
            break;
        }
    }

    for dx in -3..=3 {
        for dy in -3..=3 {
            for dz in -3..=3 {
                let t = NanoBot { x: loc.x + dx, y: loc.y + dy, z: loc.z + dz, r: 0 };
                if !best.iter().any(|&e| dist(&t, &input[e]) > input[e].r) {
                    return Some(t);
                }
            }
        }
    }

    None
}

#[aoc(day23, part2)]
pub fn solve_part2(input: &Vec<NanoBot>) -> usize {
    //println!("{:?}", input);

    let mut in_range_map : Vec<Vec<bool>> = vec![vec![false; input.len()]; input.len()];
    for i in 0..input.len() {
        for j in 0..input.len() {
            in_range_map[i][j] = dist(&input[i], &input[j]) <= input[i].r + input[j].r;
        }
    }
    let num_in_range_neighbors : Vec<usize> = in_range_map.iter().map(|row| row.iter().map(|&c| c as usize).sum()).collect();


    let mut indicies : Vec<usize> = (0..input.len()).collect();
    indicies.sort_by(|&l, &r| num_in_range_neighbors[r].cmp(&num_in_range_neighbors[l]));

    let mut best : Vec<usize> = Vec::new();
    find_most_overlapping(input, &in_range_map, &indicies, &mut Vec::new(), &mut best);
    //println!("{:?}", best);

    let matched_all = find_location(&input, &best).unwrap();

    println!("{:?}", matched_all);

    (matched_all.x.abs() + matched_all.y.abs() + matched_all.z.abs()) as usize
}