Laskettelukeskus

use std::io::stdin;
 
/// Input:
/// 5
/// 1 2
/// 1 3
/// 3 4
/// 3 5
/// 0 2 4 3 1
/// output:
/// 6
 
type RinneId = usize;
struct Polku(RinneId, RinneId);
 
#[derive(Clone, Debug)]
struct Rinne {
    #[allow(dead_code)]
    debug_id: RinneId,
    suorat_rinteet_alas: Vec<Rinne>,
    omat_aurauskerrat: usize,
}
impl Rinne {
    fn new_compute_children(id: RinneId, polut: &[Polku], aurauskerrat: &[usize]) -> Self {
        let suorat_rinteet_alas = polut
            .iter()
            .filter(|polku| polku.0 == id)
            .map(|polku| Rinne::new_compute_children(polku.1, polut, aurauskerrat))
            .collect();
        Self {
            debug_id: id,
            suorat_rinteet_alas,
            omat_aurauskerrat: aurauskerrat[id - 1],
        }
    }
    fn max_aurauskerat(&self) -> usize {
        self.omat_aurauskerrat.max(
            self.suorat_rinteet_alas
                .iter()
                .map(Rinne::max_aurauskerat)
                .sum()
        )
    }
}
 
fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut lines = stdin().lines();
    let rinteet: usize = lines.next().unwrap()?.parse()?;
    let polut: Result<Vec<Polku>, Box<dyn std::error::Error>> = lines
        .by_ref()
        .take(rinteet - 1)
        .map(|rivi| {
            let rivi = rivi?;
            // Lue luvut rivistä
            let mut luvut = rivi.split_ascii_whitespace().map(|s| s.parse());
            Ok(Polku(luvut.next().unwrap()?, luvut.next().unwrap()?))
        })
        .collect();
    let polut = polut?;
    let aurauskerrat = lines
        .next()
        .unwrap()?
        .split_ascii_whitespace()
        .map(|s| s.parse())
        .collect::<Result<Vec<usize>, _>>()?;
    assert_eq!(aurauskerrat.len(), rinteet);
 
    let tree = Rinne::new_compute_children(1, &polut, &aurauskerrat);
 
    let lumiauroja = tree.max_aurauskerat();
    println!("{lumiauroja}");
    Ok(())
}

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: TIME LIMIT EXCEEDED

Test 5

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 6

Group: 1, 2

Verdict: ACCEPTED

Test 7

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 8

Group: 1, 2

Verdict: ACCEPTED

Test 9

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 10

Group: 1, 2

Verdict: ACCEPTED

Test 11

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 12

Group: 1, 2

Verdict: ACCEPTED

Test 13

Group: 2

Verdict: TIME LIMIT EXCEEDED