Submission #1097807

---

Source Code Expand

use std::io::{self, Stdin};
use std::str::{self, FromStr};
use std::error::Error;
use std::cmp::*;
fn main() {
    let mut sc = Scanner::new();
    let n: usize = sc.ne();
    let num: Vec<i64> = (0..n).map(|_| sc.ne()).collect();
    let mut adj = vec![Vec::new(); n];
    for _ in 0..n - 1 {
        let u = sc.ne::<usize>() - 1;
        let v = sc.ne::<usize>() - 1;
        adj[u].push(v);
        adj[v].push(u);
    }
    fn dfs(adj: &Vec<Vec<usize>>, num: &Vec<i64>, cur: usize, par: usize) -> (i64, bool) {
        if adj[cur].len() == 1 {
            return (num[cur], true);
        }
        let mut kids = Vec::new();
        for nxt in &adj[cur] {
            if *nxt == par {
                continue;
            }
            let res = dfs(&adj, &num, *nxt, cur);
            if res.1 == false {
                return (0, false);
            }
            kids.push(res.0);
        }
        let sum = kids.iter().fold(0, |acc, &x| acc + x);
        if sum < num[cur] {
            return (0, false);
        }
        let ma = kids.iter().max().unwrap();
        // sum - num[cur] -> 子のペアの必要数
        if sum - num[cur] > min(sum - *ma, sum / 2) {
            return (0, false);
        }
        (2 * num[cur] - sum, true)
    }
    let ok = match n {
        2 => num[0] == num[1],
        _ => {
            let mut root = 0;
            for i in (0..n).filter(|&i| adj[i].len() > 1) {
                root = i;
                break;
            }
            let res = dfs(&adj, &num, root, root);
            // println!("{:?}", res);
            res.0 == 0 && res.1
        }
    };
    println!("{}", if ok { "YES" } else { "NO" });
}



struct Scanner {
    stdin: Stdin,
    id: usize,
    buf: Vec<u8>,
}

impl Scanner {
    fn new() -> Scanner {
        Scanner {
            stdin: io::stdin(),
            id: 0,
            buf: Vec::new(),
        }
    }
    fn next_line(&mut self) -> Option<String> {
        let mut res = String::new();
        match self.stdin.read_line(&mut res) {
            Ok(0) => return None,
            Ok(_) => Some(res),
            Err(why) => panic!("error in read_line: {}", why.description()),
        }
    }
    fn next<T: FromStr>(&mut self) -> Option<T> {
        while self.buf.len() == 0 {
            self.buf = match self.next_line() {
                Some(r) => {
                    self.id = 0;
                    r.trim().as_bytes().to_owned()
                }
                None => return None,
            };
        }
        let l = self.id;
        assert!(self.buf[l] != b' ');
        let n = self.buf.len();
        let mut r = l;
        while r < n && self.buf[r] != b' ' {
            r += 1;
        }
        let res = match str::from_utf8(&self.buf[l..r]).ok().unwrap().parse::<T>() {
            Ok(s) => Some(s),
            Err(_) => panic!("parse error"),
        };
        while r < n && self.buf[r] == b' ' {
            r += 1;
        }
        if r == n {
            self.buf.clear();
        } else {
            self.id = r;
        }
        res
    }
    fn ne<T: FromStr>(&mut self) -> T {
        self.next::<T>().unwrap()
    }
}

Submission Info

Judge Result