Skip to content

3367. Maximize Sum of Weights after Edge Removals

SourceWeekly Contest 425 Q4DifficultyHardRating2601

Description

There exists an undirected tree with n nodes numbered 0 to n - 1. You are given a 2D integer array edges of length n - 1, where edges[i] = [ui, vi, wi] indicates that there is an edge between nodes ui and vi with weight wi in the tree.

Your task is to remove zero or more edges such that:

  • Each node has an edge with at most k other nodes, where k is given.
  • The sum of the weights of the remaining edges is maximized.

Return the maximum possible sum of weights for the remaining edges after making the necessary removals.

 

Example 1:

Input: edges = [[0,1,4],[0,2,2],[2,3,12],[2,4,6]], k = 2

Output: 22

Explanation:

  • Node 2 has edges with 3 other nodes. We remove the edge [0, 2, 2], ensuring that no node has edges with more than k = 2 nodes.
  • The sum of weights is 22, and we can't achieve a greater sum. Thus, the answer is 22.

Example 2:

Input: edges = [[0,1,5],[1,2,10],[0,3,15],[3,4,20],[3,5,5],[0,6,10]], k = 3

Output: 65

Explanation:

  • Since no node has edges connecting it to more than k = 3 nodes, we don't remove any edges.
  • The sum of weights is 65. Thus, the answer is 65.

 

Constraints:

  • 2 <= n <= 105
  • 1 <= k <= n - 1
  • edges.length == n - 1
  • edges[i].length == 3
  • 0 <= edges[i][0] <= n - 1
  • 0 <= edges[i][1] <= n - 1
  • 1 <= edges[i][2] <= 106
  • The input is generated such that edges form a valid tree.

Solutions

Solution 1

Thinking

After deletions every degree is at most \(k\), and we want the maximum remaining weight. With \(n \le 10^5\) this is a tree DP.

Each subtree returns two values: the best score when the parent edge is free (up to \(k\) child edges) and when the parent edge already uses one slot (up to \(k-1\)).

A child edge is worth keeping when \(w+b-a>0\); we take the largest such gains. The root answer is the better of its two values.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int:
        def dfs(u: int, fa: int) -> Tuple[int, int]:
            s = 0
            t = []
            for v, w in g[u]:
                if v == fa:
                    continue
                a, b = dfs(v, u)
                s += a
                if (d := (w + b - a)) > 0:
                    t.append(d)
            t.sort(reverse=True)
            return s + sum(t[:k]), s + sum(t[: k - 1])

        n = len(edges) + 1
        g: List[List[Tuple[int, int]]] = [[] for _ in range(n)]
        for u, v, w in edges:
            g[u].append((v, w))
            g[v].append((u, w))
        x, y = dfs(0, -1)
        return max(x, y)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class Solution {
    private List<int[]>[] g;
    private int k;

    public long maximizeSumOfWeights(int[][] edges, int k) {
        this.k = k;
        int n = edges.length + 1;
        g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (var e : edges) {
            int u = e[0], v = e[1], w = e[2];
            g[u].add(new int[] {v, w});
            g[v].add(new int[] {u, w});
        }
        var ans = dfs(0, -1);
        return Math.max(ans[0], ans[1]);
    }

    private long[] dfs(int u, int fa) {
        long s = 0;
        List<Long> t = new ArrayList<>();
        for (var e : g[u]) {
            int v = e[0], w = e[1];
            if (v == fa) {
                continue;
            }
            var res = dfs(v, u);
            s += res[0];
            long d = w + res[1] - res[0];
            if (d > 0) {
                t.add(d);
            }
        }
        t.sort(Comparator.reverseOrder());
        for (int i = 0; i < Math.min(t.size(), k - 1); ++i) {
            s += t.get(i);
        }
        return new long[] {s + (t.size() >= k ? t.get(k - 1) : 0), s};
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
public:
    long long maximizeSumOfWeights(vector<vector<int>>& edges, int k) {
        int n = edges.size() + 1;
        vector<vector<pair<int, int>>> g(n);
        for (auto& e : edges) {
            int u = e[0], v = e[1], w = e[2];
            g[u].emplace_back(v, w);
            g[v].emplace_back(u, w);
        }
        using ll = long long;
        auto dfs = [&](this auto&& dfs, int u, int fa) -> pair<ll, ll> {
            ll s = 0;
            vector<ll> t;
            for (auto& [v, w] : g[u]) {
                if (v == fa) {
                    continue;
                }
                auto [a, b] = dfs(v, u);
                s += a;
                ll d = w + b - a;
                if (d > 0) {
                    t.push_back(d);
                }
            }
            ranges::sort(t, greater<>());
            for (int i = 0; i < min((int) t.size(), k - 1); ++i) {
                s += t[i];
            }
            return {s + (t.size() >= k ? t[k - 1] : 0), s};
        };

        auto [x, y] = dfs(0, -1);
        return max(x, y);
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
func maximizeSumOfWeights(edges [][]int, k int) int64 {
    n := len(edges) + 1
    g := make([][][]int, n)
    for _, e := range edges {
        u, v, w := e[0], e[1], e[2]
        g[u] = append(g[u], []int{v, w})
        g[v] = append(g[v], []int{u, w})
    }

    var dfs func(u, fa int) (int64, int64)
    dfs = func(u, fa int) (int64, int64) {
        var s int64
        var t []int64
        for _, e := range g[u] {
            v, w := e[0], e[1]
            if v == fa {
                continue
            }
            a, b := dfs(v, u)
            s += a
            d := int64(w) + b - a
            if d > 0 {
                t = append(t, d)
            }
        }
        sort.Slice(t, func(i, j int) bool {
            return t[i] > t[j]
        })
        for i := 0; i < min(len(t), k-1); i++ {
            s += t[i]
        }
        s2 := s
        if len(t) >= k {
            s += t[k-1]
        }
        return s, s2
    }

    x, y := dfs(0, -1)
    return max(x, y)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
function maximizeSumOfWeights(edges: number[][], k: number): number {
    const n = edges.length + 1;
    const g: [number, number][][] = Array.from({ length: n }, () => []);
    for (const [u, v, w] of edges) {
        g[u].push([v, w]);
        g[v].push([u, w]);
    }
    const dfs = (u: number, fa: number): [number, number] => {
        let s = 0;
        const t: number[] = [];

        for (const [v, w] of g[u]) {
            if (v === fa) continue;

            const [a, b] = dfs(v, u);
            s += a;
            const d = w + b - a;
            if (d > 0) t.push(d);
        }

        t.sort((a, b) => b - a);
        for (let i = 0; i < Math.min(t.length, k - 1); i++) {
            s += t[i];
        }
        const s2 = s;
        if (t.length >= k) {
            s += t[k - 1];
        }

        return [s, s2];
    };

    const [x, y] = dfs(0, -1);
    return Math.max(x, y);
}

Solution 2: Tree DP + Explicit Stack

Thinking

After the deletions every degree is at most \(k\), and the remaining weight should be as large as possible. With \(n \le 10^5\), a recursive tree DP from the root is too deep: a chain makes the call depth \(n\).

The best score of a subtree depends on whether the parent edge already uses one slot. With that slot free the node may keep up to \(k\) edges to its children; with the slot taken it may keep only \(k - 1\). Add each child at the score that does not use its own parent edge, then keep a child edge only when the gain \(w + b - a\) is positive, and only the largest of those gains.

An explicit stack of \((node, parent, state)\) runs the postorder. On entry we push the exit marker and then the children, and on exit we store the two scores. The root has no parent edge, so the answer is the larger of its two values.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution:
    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int:
        n = len(edges) + 1
        g: List[List[Tuple[int, int]]] = [[] for _ in range(n)]
        for u, v, w in edges:
            g[u].append((v, w))
            g[v].append((u, w))
        keep = [0] * n
        reserve = [0] * n
        stk = [(0, -1, 0)]
        while stk:
            u, fa, state = stk.pop()
            if state == 0:
                stk.append((u, fa, 1))
                for v, _ in g[u]:
                    if v != fa:
                        stk.append((v, u, 0))
            else:
                s = 0
                t = []
                for v, w in g[u]:
                    if v == fa:
                        continue
                    a, b = keep[v], reserve[v]
                    s += a
                    if (d := (w + b - a)) > 0:
                        t.append(d)
                t.sort(reverse=True)
                keep[u] = s + sum(t[:k])
                reserve[u] = s + sum(t[: k - 1])
        return max(keep[0], reserve[0])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class Solution {
    public long maximizeSumOfWeights(int[][] edges, int k) {
        int n = edges.length + 1;
        List<int[]>[] g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (var e : edges) {
            int u = e[0], v = e[1], w = e[2];
            g[u].add(new int[] {v, w});
            g[v].add(new int[] {u, w});
        }
        long[] keep = new long[n];
        long[] reserve = new long[n];
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int u = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push(new int[] {u, fa, 1});
                for (var e : g[u]) {
                    if (e[0] != fa) {
                        stk.push(new int[] {e[0], u, 0});
                    }
                }
            } else {
                long s = 0;
                List<Long> t = new ArrayList<>();
                for (var e : g[u]) {
                    int v = e[0], w = e[1];
                    if (v == fa) {
                        continue;
                    }
                    s += keep[v];
                    long d = w + reserve[v] - keep[v];
                    if (d > 0) {
                        t.add(d);
                    }
                }
                t.sort(Comparator.reverseOrder());
                for (int i = 0; i < Math.min(t.size(), k - 1); ++i) {
                    s += t.get(i);
                }
                reserve[u] = s;
                keep[u] = s + (t.size() >= k ? t.get(k - 1) : 0);
            }
        }
        return Math.max(keep[0], reserve[0]);
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class Solution {
public:
    long long maximizeSumOfWeights(vector<vector<int>>& edges, int k) {
        int n = edges.size() + 1;
        vector<vector<pair<int, int>>> g(n);
        for (auto& e : edges) {
            int u = e[0], v = e[1], w = e[2];
            g[u].emplace_back(v, w);
            g[v].emplace_back(u, w);
        }
        using ll = long long;
        vector<ll> keep(n), reserve(n);
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto [u, fa, state] = stk.back();
            stk.pop_back();
            if (state == 0) {
                stk.push_back({u, fa, 1});
                for (auto& [v, w] : g[u]) {
                    if (v != fa) {
                        stk.push_back({v, u, 0});
                    }
                }
            } else {
                ll s = 0;
                vector<ll> t;
                for (auto& [v, w] : g[u]) {
                    if (v == fa) {
                        continue;
                    }
                    s += keep[v];
                    ll d = w + reserve[v] - keep[v];
                    if (d > 0) {
                        t.push_back(d);
                    }
                }
                ranges::sort(t, greater<>());
                for (int i = 0; i < min((int) t.size(), k - 1); ++i) {
                    s += t[i];
                }
                reserve[u] = s;
                keep[u] = s + (t.size() >= (size_t) k ? t[k - 1] : 0);
            }
        }
        return max(keep[0], reserve[0]);
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
func maximizeSumOfWeights(edges [][]int, k int) int64 {
    n := len(edges) + 1
    g := make([][][]int, n)
    for _, e := range edges {
        u, v, w := e[0], e[1], e[2]
        g[u] = append(g[u], []int{v, w})
        g[v] = append(g[v], []int{u, w})
    }
    keep := make([]int64, n)
    reserve := make([]int64, n)
    stk := [][3]int{{0, -1, 0}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        u, fa, state := cur[0], cur[1], cur[2]
        if state == 0 {
            stk = append(stk, [3]int{u, fa, 1})
            for _, e := range g[u] {
                if e[0] != fa {
                    stk = append(stk, [3]int{e[0], u, 0})
                }
            }
        } else {
            var s int64
            var t []int64
            for _, e := range g[u] {
                v, w := e[0], e[1]
                if v == fa {
                    continue
                }
                a, b := keep[v], reserve[v]
                s += a
                d := int64(w) + b - a
                if d > 0 {
                    t = append(t, d)
                }
            }
            sort.Slice(t, func(i, j int) bool {
                return t[i] > t[j]
            })
            for i := 0; i < min(len(t), k-1); i++ {
                s += t[i]
            }
            reserve[u] = s
            if len(t) >= k {
                s += t[k-1]
            }
            keep[u] = s
        }
    }
    return max(keep[0], reserve[0])
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
function maximizeSumOfWeights(edges: number[][], k: number): number {
    const n = edges.length + 1;
    const g: [number, number][][] = Array.from({ length: n }, () => []);
    for (const [u, v, w] of edges) {
        g[u].push([v, w]);
        g[v].push([u, w]);
    }
    const keep = Array(n).fill(0);
    const reserve = Array(n).fill(0);
    const stk: [number, number, number][] = [[0, -1, 0]];
    while (stk.length) {
        const [u, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([u, fa, 1]);
            for (const [v] of g[u]) {
                if (v !== fa) {
                    stk.push([v, u, 0]);
                }
            }
        } else {
            let s = 0;
            const t: number[] = [];
            for (const [v, w] of g[u]) {
                if (v === fa) continue;
                const a = keep[v];
                const b = reserve[v];
                s += a;
                const d = w + b - a;
                if (d > 0) t.push(d);
            }
            t.sort((a, b) => b - a);
            for (let i = 0; i < Math.min(t.length, k - 1); i++) {
                s += t[i];
            }
            reserve[u] = s;
            keep[u] = s + (t.length >= k ? t[k - 1] : 0);
        }
    }
    return Math.max(keep[0], reserve[0]);
}

Comments