Laskettelukeskus

use std::{collections::HashMap, io};
 
fn main() {
    let n = input()[0];
    let mut abdict: HashMap<usize, Vec<usize>> = HashMap::new();
    for _ in 0..n - 1 {
        let ab = input();
        abdict.entry(ab[0]).or_insert(Vec::new()).push(ab[1]);
    }
    let mut needs = input();
    let mut aurat: Vec<usize> = vec![0; n];
    let mut requiredaurat: usize = 0;
    loop {
        if needs.iter().all(|&need| need == 0) {
            break;
        }
        aurat[0] += 1;
        requiredaurat += 1;
        if needs[0] > 0 {
            needs[0] -= 1;
        }
        let children = abdict
            .get(&1)
            .unwrap()
            .iter()
            .map(|&child| (child, needs[child - 1]))
            .collect::<HashMap<usize, usize>>();
        let mut nextchild = children.iter().max_by_key(|child| child.1).unwrap().0 - 1;
        let mut parent = 0;
        loop {
            aurat[parent] -= 1;
            aurat[nextchild] += 1;
            if needs[nextchild] > 0 {
                needs[nextchild] -= 1;
            }
            if abdict.contains_key(&(nextchild + 1)) {
                let newchildren = abdict
                    .get(&(nextchild + 1))
                    .unwrap()
                    .iter()
                    .map(|&child| (child, needs[child - 1]))
                    .collect::<HashMap<usize, usize>>();
                parent = nextchild.clone();
                nextchild = newchildren.iter().max_by_key(|child| child.1).unwrap().0 - 1
            } else {
                break;
            }
        }
    }
    println!("{}", requiredaurat)
}
 
fn input() -> Vec<usize> {
    let mut x = String::new();
    io::stdin().read_line(&mut x).unwrap();
    x.trim()
        .split_whitespace()
        .map(|x| x.parse::<usize>().unwrap())
        .collect()
}

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: WRONG ANSWER

Test 3

Group: 1, 2

Verdict: WRONG ANSWER

Test 4

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 5

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 6

Group: 1, 2

Verdict: ACCEPTED

Test 7

Group: 2

Verdict: ACCEPTED

Test 8

Group: 1, 2

Verdict: ACCEPTED

Test 9

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 10

Group: 1, 2

Verdict: WRONG ANSWER

Test 11

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 12

Group: 1, 2

Verdict: WRONG ANSWER

Test 13

Group: 2

Verdict: TIME LIMIT EXCEEDED