2378. Choose Edges to Maximize Score in a Tree 🔒

Description

You are given a weighted tree consisting of n nodes numbered from 0 to n - 1.

The tree is rooted at node 0 and represented with a 2D array edges of size n where edges[i] = [pari, weighti] indicates that node pari is the parent of node i, and the edge between them has a weight equal to weighti. Since the root does not have a parent, you have edges[0] = [-1, -1].

Choose some edges from the tree such that no two chosen edges are adjacent and the sum of the weights of the chosen edges is maximized.

Return the maximum sum of the chosen edges.

Note:

• You are allowed to not choose any edges in the tree, the sum of weights in this case will be 0.
• Two edges Edge1 and Edge2 in the tree are adjacent if they have a common node.
  • In other words, they are adjacent if Edge1 connects nodes a and b and Edge2 connects nodes b and c.

 

Example 1:

Input: edges = [[-1,-1],[0,5],[0,10],[2,6],[2,4]]
Output: 11
Explanation: The above diagram shows the edges that we have to choose colored in red.
The total score is 5 + 6 = 11.
It can be shown that no better score can be obtained.

Example 2:

Input: edges = [[-1,-1],[0,5],[0,-6],[0,7]]
Output: 7
Explanation: We choose the edge with weight 7.
Note that we cannot choose more than one edge because all edges are adjacent to each other.

 

Constraints:

• n == edges.length
• 1 <= n <= 105
• edges[i].length == 2
• par0 == weight0 == -1
• 0 <= pari <= n - 1 for all i >= 1.
• pari != i
• -106 <= weighti <= 106 for all i >= 1.
• edges represents a valid tree.

Solutions

Solution 1: Tree DP

We design a function \(dfs(i)\), which represents the maximum sum of the weights of selected edges in the subtree rooted at node \(i\), such that no two selected edges are adjacent. This function returns two values \((a, b)\). The first value \(a\) represents the sum of the weights of selected edges when the edge between the current node \(i\) and its parent node is selected. The second value \(b\) represents the sum of the weights of selected edges when the edge between the current node \(i\) and its parent node is not selected.

We can observe the following for the current node \(i\):

• If the edge between \(i\) and its parent node is selected, then none of the edges between \(i\) and its child nodes can be selected. In this case, the value of \(a\) for the current node is the sum of the \(b\) values of all its child nodes.
• If the edge between \(i\) and its parent node is not selected, then we can select at most one edge between \(i\) and its child nodes. In this case, the value of \(b\) for the current node is the sum of the \(a\) values of the selected child nodes and the \(b\) values of the unselected child nodes, plus the weight of the edge between \(i\) and the selected child node.

We call the function \(dfs(0)\), and the second value returned is the answer, which is the sum of the weights of selected edges when the edge between the root node and its parent node is not selected.

The time complexity is \(O(n)\), and the space complexity is \(O(n)\). Here, \(n\) is the number of nodes.

Python3JavaC++Go

class Solution:
    def maxScore(self, edges: List[List[int]]) -> int:
        def dfs(i):
            a = b = t = 0
            for j, w in g[i]:
                x, y = dfs(j)
                a += y
                b += y
                t = max(t, x - y + w)
            b += t
            return a, b

        g = defaultdict(list)
        for i, (p, w) in enumerate(edges[1:], 1):
            g[p].append((i, w))
        return dfs(0)[1]

class Solution {
    private List<int[]>[] g;

    public long maxScore(int[][] edges) {
        int n = edges.length;
        g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (int i = 1; i < n; ++i) {
            int p = edges[i][0], w = edges[i][1];
            g[p].add(new int[] {i, w});
        }
        return dfs(0)[1];
    }

    private long[] dfs(int i) {
        long a = 0, b = 0, t = 0;
        for (int[] nxt : g[i]) {
            int j = nxt[0], w = nxt[1];
            long[] s = dfs(j);
            a += s[1];
            b += s[1];
            t = Math.max(t, s[0] - s[1] + w);
        }
        b += t;
        return new long[] {a, b};
    }
}

class Solution {
public:
    long long maxScore(vector<vector<int>>& edges) {
        int n = edges.size();
        vector<vector<pair<int, int>>> g(n);
        for (int i = 1; i < n; ++i) {
            int p = edges[i][0], w = edges[i][1];
            g[p].emplace_back(i, w);
        }
        using ll = long long;
        using pll = pair<ll, ll>;
        function<pll(int)> dfs = [&](int i) -> pll {
            ll a = 0, b = 0, t = 0;
            for (auto& [j, w] : g[i]) {
                auto [x, y] = dfs(j);
                a += y;
                b += y;
                t = max(t, x - y + w);
            }
            b += t;
            return make_pair(a, b);
        };
        return dfs(0).second;
    }
};

func maxScore(edges [][]int) int64 {
    n := len(edges)
    g := make([][][2]int, n)
    for i := 1; i < n; i++ {
        p, w := edges[i][0], edges[i][1]
        g[p] = append(g[p], [2]int{i, w})
    }
    var dfs func(int) [2]int
    dfs = func(i int) [2]int {
        var a, b, t int
        for _, e := range g[i] {
            j, w := e[0], e[1]
            s := dfs(j)
            a += s[1]
            b += s[1]
            t = max(t, s[0]-s[1]+w)
        }
        b += t
        return [2]int{a, b}
    }
    return int64(dfs(0)[1])
}

Comments