#1176253 (Rust) No.417 チューリップバブル

• 問題: No.417 チューリップバブル
• ユーザー: ngtkana
• 提出日時: 2026-08-11 18:05:47
• 言語: Rust (1.94.0 + proconio + num + itertools)
• 結果: AC
• 実行時間: 1 ms / 2,000 ms
• コード長: 6,806 bytes
• コンパイル時間: 5,706 ms
• コンパイル使用メモリ: 193,464 KB
• 実行使用メモリ: 5,888 KB
• 最終ジャッジ日時: 2026-08-11 18:06:06
• 合計ジャッジ時間: 8,124 ms
• ジャッジサーバーID （参考情報）: judge1_0 / judge2_0

ファイルパターン	結果
other	AC * 40

ソースコード

use crate::flat_graph::Graph;
use proconio::input;

fn main() {
    input! {
        n: usize,
        cap: usize,
        values: [u64; n],
        edges: [(usize, usize, usize); n - 1],
    }
    let cap = cap / 2;
    let mut g = Graph::from_undirected_edges_with_weight(n, &edges);
    let (_sorted, _parent) = g.sort_undirected_tree_with_weight(0);
    let dp = dfs(0, &g, &values, vec![Some(0); cap + 1]);
    let ans = dp.last().unwrap().unwrap();
    println!("{ans}");
}

fn dfs(
    x: usize,
    g: &Graph<(usize, usize)>,
    values: &[u64],
    mut dp: Vec<Option<u64>>,
) -> Vec<Option<u64>> {
    for dp in dp.iter_mut().flatten() {
        *dp += values[x];
    }
    for &(y, w) in &g[x] {
        if dp.len() <= w {
            continue;
        }
        let mut ep = dp.to_vec();
        ep.rotate_right(w);
        ep[..w].fill(None);
        let ep = dfs(y, g, values, ep);
        for (dp, ep) in dp.iter_mut().zip(ep) {
            *dp = (*dp).max(ep);
        }
    }
    dp
}

// flat_graph {{{
// https://ngtkana.github.io/ac-adapter-rs/flat_graph/index.html
#[allow(unused_imports)]
#[allow(dead_code)]
mod flat_graph {
    use std::ops::Index;
    #[derive(Clone, Debug)]
    pub struct Graph<E> {
        start: Vec<usize>,
        tar: Vec<E>,
    }
    impl Graph<usize> {
        pub fn sort_undirected_tree(&mut self, root: usize) -> (Vec<usize>, Vec<usize>) {
            let n = self.start.len() - 1;
            assert_eq!(self.tar.len(), 2 * (n - 1));

            let mut sorted = vec![];
            let mut stack = vec![root];
            let mut parent = vec![usize::MAX; n];
            parent[root] = 0;
            while let Some(x) = stack.pop() {
                sorted.push(x);
                for &y in &self[x] {
                    if parent[y] != usize::MAX {
                        continue;
                    }
                    parent[y] = x;
                    stack.push(y);
                }
            }

            let mut i = 0;
            let mut j = 0;
            for x in 0..n {
                while j < self.start[x + 1] {
                    if self.tar[j] != parent[x] {
                        self.tar.swap(i, j);
                        i += 1;
                    }
                    j += 1;
                }
                self.start[x + 1] = i;
            }
            assert_eq!(i, n - 1);

            self.tar.truncate(n - 1);
            (sorted, parent)
        }
        pub fn from_directed_edges(n: usize, edges: &[(usize, usize)]) -> Self {
            Self::from_edges_generic(
                n,
                edges.len(),
                edges.iter().map(|&(i, _)| i),
                edges.iter().map(|&(i, j)| (i, j)),
            )
        }
        pub fn from_undirected_edges(n: usize, edges: &[(usize, usize)]) -> Self {
            Self::from_edges_generic(
                n,
                edges.len() * 2,
                edges.iter().flat_map(|&(i, j)| [i, j]),
                edges.iter().flat_map(|&(i, j)| [(i, j), (j, i)]),
            )
        }
    }
    impl<T: Copy + Default> Graph<(usize, T)> {
        pub fn from_directed_edges_with_weight(n: usize, edges: &[(usize, usize, T)]) -> Self {
            Self::from_edges_generic(
                n,
                edges.len(),
                edges.iter().map(|&(i, _, _)| i),
                edges.iter().map(|&(i, j, w)| (i, (j, w))),
            )
        }
        pub fn from_undirected_edges_with_weight(n: usize, edges: &[(usize, usize, T)]) -> Self {
            Self::from_edges_generic(
                n,
                edges.len() * 2,
                edges.iter().flat_map(|&(i, j, _)| [i, j]),
                edges
                    .iter()
                    .flat_map(|&(i, j, w)| [(i, (j, w)), (j, (i, w))]),
            )
        }
        pub fn sort_undirected_tree_with_weight(
            &mut self,
            root: usize,
        ) -> (Vec<usize>, Vec<usize>) {
            let n = self.start.len() - 1;
            assert_eq!(self.tar.len(), 2 * (n - 1));

            let mut sorted = vec![];
            let mut stack = vec![root];
            let mut parent = vec![usize::MAX; n];
            parent[root] = 0;
            while let Some(x) = stack.pop() {
                sorted.push(x);
                for &(y, _) in &self[x] {
                    if parent[y] != usize::MAX {
                        continue;
                    }
                    parent[y] = x;
                    stack.push(y);
                }
            }

            let mut i = 0;
            let mut j = 0;
            for x in 0..n {
                while j < self.start[x + 1] {
                    if self.tar[j].0 != parent[x] {
                        self.tar.swap(i, j);
                        i += 1;
                    }
                    j += 1;
                }
                self.start[x + 1] = i;
            }
            assert_eq!(i, n - 1);

            self.tar.truncate(n - 1);
            (sorted, parent)
        }
    }
    impl<E: Default + Clone> Graph<E> {
        fn from_edges_generic(
            n: usize,
            m: usize,
            src: impl Iterator<Item = usize>,
            edges: impl Iterator<Item = (usize, E)>,
        ) -> Self {
            let mut start = vec![0; n + 1];
            for i in src {
                start[i + 1] += 1;
            }
            for i in 0..n {
                start[i + 1] += start[i];
            }
            let edge_count = m;
            let mut tar = vec![E::default(); edge_count];
            for (i, e) in edges {
                tar[start[i]] = e;
                start[i] += 1;
            }
            start.rotate_right(1);
            start[0] = 0;
            Self { start, tar }
        }
    }
    impl<E> Graph<E> {
        pub fn iter(&self) -> Iter<'_, E> {
            Iter {
                index: 0,
                graph: self,
            }
        }
    }
    impl<E> Index<usize> for Graph<E> {
        type Output = [E];
        fn index(&self, index: usize) -> &Self::Output {
            &self.tar[self.start[index]..self.start[index + 1]]
        }
    }
    impl<'a, E> IntoIterator for &'a Graph<E> {
        type Item = &'a [E];
        type IntoIter = Iter<'a, E>;
        fn into_iter(self) -> Self::IntoIter {
            self.iter()
        }
    }
    pub struct Iter<'a, E> {
        index: usize,
        graph: &'a Graph<E>,
    }
    impl<'a, E> Iterator for Iter<'a, E> {
        type Item = &'a [E];
        fn next(&mut self) -> Option<Self::Item> {
            if self.index + 1 == self.graph.start.len() {
                None
            } else {
                self.index += 1;
                Some(&self.graph[self.index - 1])
            }
        }
    }
}
// }}}