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] = [par_i, weight_i] indicates that node par_i is the parent of node i, and the edge between them has a weight equal to weight_i. 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 Edge₁ and Edge₂ in the tree are adjacent if they have a common node.
- In other words, they are adjacent if Edge₁ connects nodes a and b and Edge₂ 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 <= 10⁵
edges[i].length == 2
par₀ == weight₀ == -1
0 <= par_i <= n - 1 for all i >= 1.
par_i != i
-10⁶ <= weight_i <= 10⁶ for all i >= 1.
edges represents a valid tree.

Solutions

Solution 2: Tree DP + Explicit Stack

Thinking

Choose a set of non-adjacent tree edges with maximum weight. With $n \le 10^5$, deciding each edge by recursion from the root is too deep: a chain makes the call depth $n$.

Each node needs two values: the best weight $a$ when the edge to its parent is taken, and the best weight $b$ when that edge is not taken. Taking the parent edge forbids every edge to a child, so $a$ is the sum of the children’s $b$ values. Leaving it free allows at most one child edge, and the gain of switching that child is its $a$ minus its $b$ plus the edge weight.

An explicit stack of $(node, state)$ runs the postorder. On entry we push the exit marker and then the children, and on exit we store $a$ and $b$. The root has no parent edge, so the answer is the root’s $b$.

For each node $i$ we compute two values $(a, b)$. The value $a$ is the maximum weight of selected edges in the subtree when the edge between $i$ and its parent is selected. The value $b$ is that weight when the edge is not selected. An explicit stack walks from the root in postorder and writes the pair when the node is left.

For the current node $i$:

If the edge between $i$ and its parent is selected, then none of the edges between $i$ and its children can be selected. In this case $a$ is the sum of the $b$ values of all children.
If that edge is not selected, then at most one edge from $i$ to a child can be selected. In this case $b$ is the selected child’s $a$, plus the $b$ values of the other children, plus the weight of the selected edge. A negative gain leaves the edge unselected.

The root has no parent edge, so the answer is the root’s $b$.

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

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};
    }
}


// Solution 2
class Solution {
    public long maxScore(int[][] edges) {
        int n = edges.length;
        List<int[]>[] 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});
        }
        long[][] down = new long[n][2];
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], state = cur[1];
            if (state == 0) {
                stk.push(new int[] {i, 1});
                for (int[] nxt : g[i]) {
                    stk.push(new int[] {nxt[0], 0});
                }
            } else {
                long a = 0, b = 0, t = 0;
                for (int[] nxt : g[i]) {
                    int j = nxt[0], w = nxt[1];
                    long x = down[j][0], y = down[j][1];
                    a += y;
                    b += y;
                    t = Math.max(t, x - y + w);
                }
                b += t;
                down[i][0] = a;
                down[i][1] = b;
            }
        }
        return down[0][1];
    }
}

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;
    }
};


// Solution 2
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;
        vector<pair<ll, ll>> down(n);
        vector<array<int, 2>> stk{{0, 0}};
        while (!stk.empty()) {
            auto [i, state] = stk.back();
            stk.pop_back();
            if (state == 0) {
                stk.push_back({i, 1});
                for (auto& [j, w] : g[i]) {
                    stk.push_back({j, 0});
                }
            } else {
                ll a = 0, b = 0, t = 0;
                for (auto& [j, w] : g[i]) {
                    auto [x, y] = down[j];
                    a += y;
                    b += y;
                    t = max(t, x - y + w);
                }
                b += t;
                down[i] = {a, b};
            }
        }
        return down[0].second;
    }
};

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]



# Solution 2
class Solution:
    def maxScore(self, edges: List[List[int]]) -> int:
        n = len(edges)
        g = [[] for _ in range(n)]
        for i, (p, w) in enumerate(edges[1:], 1):
            g[p].append((i, w))
        down = [(0, 0)] * n
        stk = [(0, 0)]
        while stk:
            i, state = stk.pop()
            if state == 0:
                stk.append((i, 1))
                for j, _ in g[i]:
                    stk.append((j, 0))
            else:
                a = b = t = 0
                for j, w in g[i]:
                    x, y = down[j]
                    a += y
                    b += y
                    t = max(t, x - y + w)
                b += t
                down[i] = (a, b)
        return down[0][1]

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])
}


// Solution 2
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})
	}
	down := make([][2]int, n)
	stk := [][2]int{{0, 0}}
	for len(stk) > 0 {
		cur := stk[len(stk)-1]
		stk = stk[:len(stk)-1]
		i, state := cur[0], cur[1]
		if state == 0 {
			stk = append(stk, [2]int{i, 1})
			for _, e := range g[i] {
				stk = append(stk, [2]int{e[0], 0})
			}
		} else {
			var a, b, t int
			for _, e := range g[i] {
				j, w := e[0], e[1]
				x, y := down[j][0], down[j][1]
				a += y
				b += y
				t = max(t, x-y+w)
			}
			b += t
			down[i] = [2]int{a, b}
		}
	}
	return int64(down[0][1])
}

2378 - Choose Edges to Maximize Score in a Tree

2378. Choose Edges to Maximize Score in a Tree

Description

Solutions

Solution 2: Tree DP + Explicit Stack

All Problems

All Solutions

Welcome to Subscribe On Youtube

2378. Choose Edges to Maximize Score in a Tree

Description

Solutions

Solution 2: Tree DP + Explicit Stack

All Problems

All Solutions