codejam treasure problem

treasure.rs

/**
* https://codingcompetitions.withgoogle.com/codejam/round/0000000000432dfc/0000000000433101
* @robertkingnz
**/

use std::collections::{HashMap, HashSet};

fn get_line() -> String {
    let mut buf = String::new();
    std::io::stdin().read_line(&mut buf).unwrap();
    buf.trim().into()
}

fn get_usize() -> usize {
    get_line().parse::<usize>().unwrap()
}

fn get_usizes() -> Vec<usize> {
    let mut v = vec![];
    for part in get_line().split(" ") {
        v.push(part.parse::<usize>().unwrap())
    }
    v
}

#[derive(Debug)]
struct Chest {
    lock: usize,
    keys: HashMap<usize, usize>,
    is_open: bool,
    order: Option<usize>
}

impl Chest {
    fn read() -> Self {
        let nums = get_usizes();
        let mut keys: HashMap<usize,usize> = HashMap::new();
        for &k in nums.iter().skip(2) {
            *keys.entry(k).or_default() += 1;
        }
        return Chest {
            lock: nums[0],
            keys,
            is_open: false,
            order: None
        }
    }
}

fn can_open(i: usize, chests: &Vec<Chest>, my_keys: &HashMap<usize, usize>) -> bool {
    let mut keys = my_keys.clone();
    *keys.entry(chests[i].lock).or_default() -= 1;
    for (k, count) in &chests[i].keys {
        *keys.entry(*k).or_default() += *count;

    }
    let mut graph: HashMap<usize, HashSet<usize>> = HashMap::new();
    for i2 in 0..chests.len() {
        if i == i2 {
            continue;
        }
        let chest = &chests[i2];
        if chest.is_open {
            continue;
        }
        graph.entry(chest.lock)
            .or_default()
            .extend(chest.keys.keys());
    }
    let mut stack = vec![];
    let mut visited: HashSet<usize> = HashSet::new();
    for (k, count) in keys {
        if count > 0 {
            stack.push(k);
            visited.insert(k);
        }
    }

    while let Some(k) = stack.pop() {
        if let Some(children) = graph.get(&k) {
            for &child in children {
                if !visited.contains(&child) {
                    visited.insert(child);
                    stack.push(child);
                }
            }
        }
    }
    for i2 in 0..chests.len() {
        if i == i2 || chests[i2].is_open {
            continue;
        }
        if !visited.contains(&chests[i2].lock) {
            return false;
        }
    }
    true
}

fn solve_test_case() -> Option<String> {
    let problem_info = get_usizes();
    let mut my_keys: HashMap<usize,usize> = HashMap::new();
    let mut total_lock_count: HashMap<usize,usize> = HashMap::new();
    for k in get_usizes() {
        *my_keys.entry(k).or_default() += 1;
    }
    let mut total_key_count = my_keys.clone();
    let n = problem_info[1];
    let mut chests: Vec<Chest> = Vec::with_capacity(n);
    for _ in 0..n {
        let chest = Chest::read();
        for (key, count) in &chest.keys {
            *total_key_count.entry(*key).or_default() += *count;
        }
        *total_lock_count.entry(chest.lock).or_default() += 1;
        chests.push(chest);
    }
    for (lock, lock_count) in total_lock_count {
        if lock_count > *total_key_count.get(&lock).unwrap_or(&0) {
            return None;
        }
    }
    let mut i = 0;
    let mut result = vec![];
    while i < n {
        let has_key = *my_keys.get(&chests[i].lock).unwrap_or(&0) > 0;
        if !chests[i].is_open && has_key && can_open(i, &chests, &my_keys) {
            for (&k, &count) in &chests[i].keys {
                *my_keys.entry(k).or_default() += count;
            }
            *my_keys.entry(chests[i].lock).or_default() -= 1;
            chests[i].is_open = true;
            result.push((i+1).to_string());
            i = 0;
        } else {
            i += 1;
        }
    }
    if result.len() != n {
        return None;
    }
    result.join(" ").into()
}

fn main() {
    let t = get_usize();
    for case_num in 1..=t {
        if let Some(ans) = solve_test_case() {
            println!("Case #{}: {}", case_num, ans);
        } else {
            println!("Case #{}: IMPOSSIBLE", case_num);
        }
    }
}