LeetCode每日一题(2538. Difference Between Maximum and Minimum Price Sum)

There exists an undirected and initially unrooted tree with n nodes indexed from 0 to n - 1. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Each node has an associated price. You are given an integer array price, where price[i] is the price of the ith node.

The price sum of a given path is the sum of the prices of all nodes lying on that path.

The tree can be rooted at any node root of your choice. The incurred cost after choosing root is the difference between the maximum and minimum price sum amongst all paths starting at root.

Return the maximum possible cost amongst all possible root choices.

Example 1:

Input: n = 6, edges = [[0,1],[1,2],[1,3],[3,4],[3,5]], price = [9,8,7,6,10,5]
Output: 24

Explanation: The diagram above denotes the tree after rooting it at node 2. The first part (colored in red) shows the path with the maximum price sum. The second part (colored in blue) shows the path with the minimum price sum.

• The first path contains nodes [2,1,3,4]: the prices are [7,8,6,10], and the sum of the prices is 31.
• The second path contains the node [2] with the price [7].
  The difference between the maximum and minimum price sum is 24. It can be proved that 24 is the maximum cost.

Example 2:

Input: n = 3, edges = [[0,1],[1,2]], price = [1,1,1]
Output: 2

Explanation: The diagram above denotes the tree after rooting it at node 0. The first part (colored in red) shows the path with the maximum price sum. The second part (colored in blue) shows the path with the minimum price sum.

• The first path contains nodes [0,1,2]: the prices are [1,1,1], and the sum of the prices is 3.
• The second path contains node [0] with a price [1].
  The difference between the maximum and minimum price sum is 2. It can be proved that 2 is the maximum cost.

Constraints:

• 1 <= n <= 105
• edges.length == n - 1
• 0 <= ai, bi <= n - 1
• edges represents a valid tree.
• price.length == n
• 1 <= price[i] <= 105

---

这题之所以被归为 hard很大一部分原因是因为 test case

有最大差值的 path 一定是 leaf node 到 leaf node 的假设 leaf_a 节点到 leaf_b 节点的最大 sum 是 s, 那最大差值就是 max(s - price[leaf_a], s - price[leaf_b])。

假设 to_leaf_sums[i]是 node i 到所有 leaf node 的 sum 的最大值 not_to_leaf_sums[i]是 node i 到所有 leaf node 的 father node 的 sum 的最大值, 那经过 node i 的最大 diff 就是 to_leaf_sums[i] + not_to_leaf_sums[i], 这里要注意 两个值所对应的 leaf node 不能是同一个 leaf node。

假设 children[i]是 node i 的所有 child nodes, 即 children[i] = [child1, child2, …]

to_leaf_sums[i] = max(to_leaf_sums[children[i][0]], to_leaf_sums[children[i][1]], …, to_leaf_sums[children[i][last]]) + price[i]

not_to_leaf_sums 同理如上

这样我们按任意一个 node 作为 root 进行计算即可。

之所以说这题归为 hard 的原因是 test case 是因为按上面的解法会导致超时 开始我以为是自己的缓存写的有问题但是试着改过几遍之后发现好像不是缓存引起的。再看超时的那个 tast case 发现是一个大的星型 graph 一共 50000 个节点 节点 0 与其他 49999 个节点直接相连。因为开始没有考虑这种情况 所以在计算 to_leaf_sums[child1] + to_leaf_sums[child2] + price[curr_node]的时候child1 和 child2 是通过两层遍历来进行的 这就导致了这里的时间复杂度变成了 O(n²), 所以会导致超时。解决办法倒是不复杂 将 child node 的 to_leaf_sums 和 not_to_leaf_sums 分别进行排序 计算的时候分别取最大值 如果两个值是来自于相同的 leaf node 的 那就分别再取两个最大值

---

use std::collections::{BinaryHeap, HashMap};

impl Solution {
    fn dp(conn_set: &Vec<Vec<usize>>, prices: &Vec<i64>, parent: usize, node: usize, ans: &mut i64, cache: &mut HashMap<usize, (i64, i64)>) -> (i64, i64) {
        let conns = &conn_set[node];
        // 如果是leaf node则直接返回
        if conns.len() == 1 && conns[0] == parent {
            return (prices[node], 0);
        }
        if let Some(c) = cache.get(&node) {
            return *c;
        }
        let mut to_leaf_sums = BinaryHeap::new();
        let mut not_to_leaf_sums = BinaryHeap::new();
        for &conn in conns {
            if conn == parent {
                continue;
            }
            // child node到leaf node与parent of leaf node的sum
            let (to_leaf, not_to_leaf) = Solution::dp(conn_set, prices, node, conn, ans, cache);
            to_leaf_sums.push((to_leaf, conn));
            not_to_leaf_sums.push((not_to_leaf, conn));
        }
        if to_leaf_sums.is_empty() {
            return (0, 0);
        }
        // 只有一个子节点的情况
        if to_leaf_sums.len() == 1 {
            let (to_leaf, _) = to_leaf_sums.pop().unwrap();
            let (not_to_leaf, _) = not_to_leaf_sums.pop().unwrap();
            *ans = (*ans).max(to_leaf.max(not_to_leaf + prices[node]));
            cache.insert(node, (to_leaf, not_to_leaf));
            return (to_leaf + prices[node], not_to_leaf + prices[node]);
        }
        // 到leaf node的最大值
        let (to_leaf, to_leaf_conn) = to_leaf_sums.pop().unwrap();
        // 到parent of leaf node的最大值
        let (not_to_leaf, not_to_leaf_conn) = not_to_leaf_sums.pop().unwrap();
        // 两个最大值来自同一节点的情况
        if to_leaf_conn == not_to_leaf_conn {
            let (next_to_leaf, _) = to_leaf_sums.pop().unwrap();
            let (next_not_to_leaf, _) = not_to_leaf_sums.pop().unwrap();
            *ans = (*ans).max((to_leaf + next_not_to_leaf + prices[node]).max(not_to_leaf + next_to_leaf + prices[node]));
            let res = (to_leaf + prices[node], not_to_leaf + prices[node]);
            cache.insert(node, res);
            return res;
        }
        // 两个最大值来自不同节点的情况
        *ans = (*ans).max(to_leaf + not_to_leaf + prices[node]);
        let res = (to_leaf + prices[node], not_to_leaf + prices[node]);
        cache.insert(node, res);
        res
    }
    pub fn max_output(n: i32, edges: Vec<Vec<i32>>, price: Vec<i32>) -> i64 {
        let mut conn_set = vec![Vec::new(); n as usize];
        for edge in edges {
            conn_set[edge[0] as usize].push(edge[1] as usize);
            conn_set[edge[1] as usize].push(edge[0] as usize);
        }
        let prices: Vec<i64> = price.into_iter().map(|v| v as i64).collect();
        let mut ans = 0;
        Solution::dp(&conn_set, &prices, 100001, 0, &mut ans, &mut HashMap::new());
        ans
    }
}