2024_day5.rs

use std::collections::HashMap;
use std::collections::HashSet;
use std::cmp::Ordering;
use std::env;
use std::fs;

fn get_arg() -> String {
    let opt_filename = env::args().nth(1);
    match opt_filename {
        Some(f) => f,
        None => "example".into(),
    }
}

fn get_content(filename: String) -> String {
    fs::read_to_string(filename)
        .unwrap()
        .trim()
        .to_string()
}

fn split_lines(s: String) -> Vec<String> {
    s.split('\n')
        .map(|line| line.to_string())
        .collect()
}


fn main() {
    let lines = split_lines(get_content(get_arg()));
    println!("Part 1: {}", part1(&lines));
    println!("Part 2: {}", part2(&lines));
}

fn part2(lines: &[String]) -> i64 {
    let (orderings, updates) = parse_lines(lines);

    let mut ordering_map: HashMap<i64, HashSet<i64>> = HashMap::new();
    for (l, r) in orderings.iter() {
        if !ordering_map.contains_key(l) {
            ordering_map.insert(*l, HashSet::new());
        }

        ordering_map.get_mut(l).unwrap().insert(*r);
    }

    let mut total = 0;
    for u in updates.iter() {
        if !is_ordered(&orderings, u) {
            let fixed = fix_ordering(&ordering_map, u);
            total += fixed[fixed.len() / 2];
        }
    }
    total
}

fn fix_ordering(ordering_map: &HashMap<i64, HashSet<i64>>, u: &[i64]) -> Vec<i64> {
    let mut copy: Vec<i64> = u.iter().map(|i| *i).collect();
    copy.sort_by(|a, b| {
        let a_before = ordering_map.get(a);
        let b_before = ordering_map.get(b);

        if let Some(before) = a_before {
            if before.contains(b) {
                return Ordering::Less;
            }
        }

        if let Some(before) = b_before {
            if before.contains(a) {
                return Ordering::Greater;
            }
        }

        Ordering::Equal
    });
    copy
}

fn part1(lines: &[String]) -> i64 {
    let (orderings, updates) = parse_lines(lines);

    let mut total = 0;
    for u in updates.iter() {
        if is_ordered(&orderings, u) {
            total += u[u.len() / 2];
        }
    }
    total
}

fn is_ordered(orderings: &[(i64, i64)], u: &[i64]) -> bool {
    let mut positions = HashMap::new();
    for (i, page) in u.iter().enumerate() {
        positions.insert(page, i);
    }


    for (first, second) in orderings.iter() {
        match (positions.get(first), positions.get(second)) {
            (Some(f), Some(s)) => {
                if f >= s {
                    return false;
                }
            },
            _ => {},
        }
    }

    true
}

fn parse_lines(lines: &[String]) -> (Vec<(i64, i64)>, Vec<Vec<i64>>) {
    let mut orderings = Vec::new();
    let mut updates = Vec::new();

    let mut lines_iter = lines.iter();
    loop {
        let option_line = lines_iter.next();
        if let Some(line) = option_line {
            if line.trim().is_empty() {
                break;
            }
            orderings.push(parse_ordering(line));
        } else {
            break;
        }
    }

    loop {
        let option_line = lines_iter.next();
        if let Some(line) = option_line {
            updates.push(parse_update(line));
        } else {
            break;
        }
    }

    (orderings, updates)
}

fn parse_ordering(line: &str) -> (i64, i64) {
    let parts: Vec<&str> = line.split('|').collect();
    (parts[0].parse().unwrap(), parts[1].parse().unwrap())
}

fn parse_update(line: &str) -> Vec<i64> {
    line.split(',').map(|s| s.parse().unwrap()).collect()
}