Laskettelukeskus

use std::cmp::max;
use std::collections::{HashMap, HashSet};
use std::io;

fn main() {
    let stdin = io::stdin();

    let slope_count = { // AKA n
        let mut input: String = String::new();
        stdin.read_line(&mut input).unwrap();
        input.replace("\r", "").replace("\n", "").parse::<usize>().unwrap()
    };

    // All slopes indexed by their IDs.
    let mut slopes = {
        let mut slopes: Vec<SkiSlope> = Vec::with_capacity(slope_count);
        for _ in 0..slope_count {
            slopes.push(SkiSlope {
                upward_connections: HashSet::new(),
                downward_connections: HashSet::new(),
                plows: 0, // Will be initialized later.
                child_plows: 0 // Will be initialized later.
            });
        }
        slopes
    };

    // Read routes
    {
        // Read line slope_count - 1 times
        for _ in 1..slope_count {
            let mut input: String = String::new();
            stdin.read_line(&mut input).unwrap();

            let route: Vec<usize> = input.replace("\r", "").replace("\n", "").split_whitespace().map(|value| {
                value.parse::<usize>().unwrap() - 1 // Indexes are given starting from 1, so they need to be decremented
            }).collect();

            // A route exists from route[0] (upper) to route[1] (lower).

            slopes[route[1]].upward_connections.insert(route[0]);
            slopes[route[0]].downward_connections.insert(route[1]);
        }
    };

    // Read number of required plows.
    {
        let mut input: String = String::new();
        stdin.read_line(&mut input).unwrap();

        input
            .replace("\r", "")
            .replace("\n", "")
            .split_whitespace()
            .enumerate()
            .for_each(|(i, value)| {
                slopes[i].plows = value.parse().unwrap();
            });
    }

    // List of ski slopes indexed by their depth from the root slope. Depth is 0 for the root slope.
    // Otherwise this could be calculated while taking input, but it's possible depth would be determined before the parent's depth is calculated.
    let depth_map: Vec<HashSet<usize>> = {
        let mut depth_map: Vec<HashSet<usize>> = Vec::new();
        depth_map.insert(0, HashSet::from([ 0 ])); // Root slope is at depth 0.

        // Calculate required plows from bottom to top.

        // key-value pairs of slopes to be processed
        let mut loop_slopes: HashMap<usize, usize> = HashMap::with_capacity(1);
        loop_slopes.insert(0, 0); // Adds root slope manually.
        while !&loop_slopes.is_empty() {
            let mut next_slopes = HashMap::new();
            for (parent_slope_id, depth) in &loop_slopes {
                let slope = &slopes[*parent_slope_id];

                let new_depth = depth + 1;
                if depth_map.len() >= new_depth {
                    depth_map.push(HashSet::new());
                }

                let set = &mut depth_map[new_depth];
                for lower_slope_id in slope.downward_connections.clone() {
                    set.insert(lower_slope_id);
                    next_slopes.insert(lower_slope_id, new_depth);
                }
            }
            loop_slopes = next_slopes;
        }

        depth_map
    };

    // Calculate required plows in reverse depth order.
    for slope_ids in depth_map.iter().rev() {
        for slope_id in slope_ids {
            let slope = &mut slopes[*slope_id];
            slope.plows = max(slope.plows, slope.child_plows);
            let plows = slope.plows;
            for upper_slope_id in slope.upward_connections.clone() {
                let mut upper_slope = &mut slopes[upper_slope_id];
                upper_slope.child_plows += plows;
            }
        }
    }

    println!("{}", slopes[0].plows); // Print required plows for the root ski slope.
}

struct SkiSlope {
    upward_connections: HashSet<usize>,
    downward_connections: HashSet<usize>,
    plows: u64, // number of required plows,
    child_plows: u64 // required plows of all children combined
}

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: ACCEPTED

Test 3

Group: 1, 2

Verdict: ACCEPTED

Test 4

Group: 2

Verdict: ACCEPTED

Test 5

Group: 2

Verdict: ACCEPTED

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: ACCEPTED

Test 10

Group: 1, 2

Verdict: ACCEPTED

Test 11

Group: 2

Verdict: ACCEPTED

Test 12

Group: 1, 2

Verdict: ACCEPTED

Test 13

Group: 2

Verdict: ACCEPTED