extern crate smallvec; // 0.6.3

trait ToZeroBased {
    fn to_zero_based(usize) -> usize;
    fn to_zero_based_pair(pair: (usize, usize)) -> (usize, usize) {
        let (a, b) = pair;
        (Self::to_zero_based(a), Self::to_zero_based(b))
    }
}

struct ZeroBased;
impl ToZeroBased for ZeroBased {
    fn to_zero_based(i: usize) -> usize { i }
}

struct OneBased;
impl ToZeroBased for OneBased {
    fn to_zero_based(i: usize) -> usize { i - 1 }
}

use std::collections::HashMap;
use smallvec::SmallVec;

type Weight = i32;
const THRESHOLD: usize = 8;
type Graph = Vec<SmallVec<[usize; THRESHOLD]>>;

fn tree_sum<I: ToZeroBased>(weights: &[Weight], edges: &[(usize, usize)]) -> Weight {
    let answers = &mut HashMap::new();
    let mut graph = vec![SmallVec::new(); weights.len()];

    for &edge in edges {
        let (from, to) = I::to_zero_based_pair(edge);

        graph[from].push(to);
        graph[to].push(from);
    }

    tree_sum_recurrent(answers, weights, &graph, 0, 0)
}

fn tree_sum_unilecs(weights: &[Weight], edges: &[(usize, usize)]) -> Weight {
    tree_sum::<OneBased>(weights, edges)
}

fn tree_sum_recurrent(
    answers: &mut HashMap<usize, Weight>,
    weights: &[Weight],
    graph: &Graph,
    prev: usize,
    center: usize,
) -> Weight {
    if !answers.contains_key(&center) {
        let neighbors = graph[center].iter();

        let without_center = neighbors.clone()
            .filter(|&&i| i != prev)
            .map(|&i| tree_sum_recurrent(answers, weights, graph, center, i))
            .sum();

        let mut without_neighbors = 0;
        for &curr in neighbors.filter(|&&i| i != prev) {
            without_neighbors += graph[curr]
                .iter()
                .filter(|&&i| i != center)
                .map(|&i| tree_sum_recurrent(answers, weights, graph, curr, i))
                .sum::<Weight>();
        }

        let with_center = weights[center] + without_neighbors;

        answers.insert(center, std::cmp::max(with_center, without_center));
    }

    answers[&center]
}

fn main() {
    let weights = [2, 1, 1, 1];
    let edges = [(0, 1), (0, 2), (0, 3)];
    println!("{}", tree_sum::<ZeroBased>(&weights, &edges)); // 3

    let weights = [1, 1, 0, 1];
    let edges = [(0, 1), (0, 2), (1, 3)];
    println!("{}", tree_sum::<ZeroBased>(&weights, &edges)); // 2
}

#[test]
fn unilecs_test_case() {
    let weights = [1, 1, 0, 1];
    let edges = [(1, 2), (1, 3), (2, 4)];

    assert_eq!(tree_sum_unilecs(&weights, &edges), 2);
}

#[test]
fn greedy_algorithm_countercase() {
    /*
     *      1
     *      |
     * 1 -- 2 -- 1
     */

    let weights = [2, 1, 1, 1];
    let edges = [(0, 1), (0, 2), (0, 3)];

    // Naive greedy algorithm would return wrong result `2`,
    // while the right answer in this case is `3`
    assert_eq!(tree_sum::<ZeroBased>(&weights, &edges), 3);
}

#[test]
fn returns_zero_for_all_negative_weights() {
    let weights = [-1, -1, -1];
    let edges = [(0, 1), (1, 2)];

    assert_eq!(tree_sum::<ZeroBased>(&weights, &edges), 0);
}

#[test]
fn returns_positive_weight_for_the_only_positive_weight() {
    let weights = [-1, -1, -1, 1];
    let edges = [(0, 1), (1, 2), (2, 3)];

    assert_eq!(tree_sum::<ZeroBased>(&weights, &edges), 1);
}

#[test]
fn stress_test() {
    const LEN: usize = 100;
    let weights = vec![1; LEN];
    let edges = (1..LEN).map(|i| (i - 1, i)).collect::<Vec<_>>();

    assert_eq!(tree_sum::<ZeroBased>(&weights, &edges), (LEN as Weight) / 2);
}