Create Components With Same Value

hard tree dfs math enumeration

Problem

An undirected tree has n nodes with values nums[i] and an edge list. Delete some edges so the tree splits into components that all have the same value sum. Return the maximum number of edges you can delete.

Inputnums = [6,2,2,2,6], edges = [[0,1],[1,2],[1,3],[3,4]]

Output2

Total is 18. Cutting into 3 components each summing to 6 ({0},{1,2,3},{4}) removes 2 edges.

nums array

edges as [u,v] list

def components_with_same_value(nums, edges):
    n = len(nums)
    g = [[] for _ in range(n)]
    for a, b in edges:
        g[a].append(b); g[b].append(a)
    total = sum(nums)
    best_max = max(nums)
    def can(target):
        ok = [True]
        def dfs(u, p):
            s = nums[u]
            for v in g[u]:
                if v != p:
                    s += dfs(v, u)
            if s == target:
                return 0
            if s > target:
                ok[0] = False
            return s
        return dfs(0, -1) == 0 and ok[0]
    for parts in range(n, 0, -1):
        if total % parts: continue
        target = total // parts
        if target < best_max: continue
        if can(target):
            return parts - 1
    return 0

function componentsWithSameValue(nums, edges) {
  const n = nums.length;
  const g = Array.from({length: n}, () => []);
  for (const [a, b] of edges) { g[a].push(b); g[b].push(a); }
  const total = nums.reduce((x, y) => x + y, 0);
  const bestMax = Math.max(...nums);
  function can(target) {
    let ok = true;
    function dfs(u, p) {
      let s = nums[u];
      for (const v of g[u]) if (v !== p) s += dfs(v, u);
      if (s === target) return 0;
      if (s > target) ok = false;
      return s;
    }
    return dfs(0, -1) === 0 && ok;
  }
  for (let parts = n; parts >= 1; parts--) {
    if (total % parts) continue;
    const target = total / parts;
    if (target < bestMax) continue;
    if (can(target)) return parts - 1;
  }
  return 0;
}

import java.util.*;
class Solution {
    List<Integer>[] g; int[] nums; int target; boolean ok;
    public int componentValueSum(int[] nums, int[][] edges) {
        int n = nums.length; this.nums = nums;
        g = new List[n];
        for (int i = 0; i < n; i++) g[i] = new ArrayList<>();
        for (int[] e : edges) { g[e[0]].add(e[1]); g[e[1]].add(e[0]); }
        int total = 0, bestMax = 0;
        for (int v : nums) { total += v; bestMax = Math.max(bestMax, v); }
        for (int parts = n; parts >= 1; parts--) {
            if (total % parts != 0) continue;
            target = total / parts;
            if (target < bestMax) continue;
            ok = true;
            if (dfs(0, -1) == 0 && ok) return parts - 1;
        }
        return 0;
    }
    int dfs(int u, int p) {
        int s = nums[u];
        for (int v : g[u]) if (v != p) s += dfs(v, u);
        if (s == target) return 0;
        if (s > target) ok = false;
        return s;
    }
}

class Solution {
public:
    vector<vector<int>> g; vector<int> nums; int target; bool ok;
    int dfs(int u, int p) {
        int s = nums[u];
        for (int v : g[u]) if (v != p) s += dfs(v, u);
        if (s == target) return 0;
        if (s > target) ok = false;
        return s;
    }
    int componentValueSum(vector<int>& nums_, vector<vector<int>>& edges) {
        int n = nums_.size(); nums = nums_;
        g.assign(n, {});
        for (auto& e : edges) { g[e[0]].push_back(e[1]); g[e[1]].push_back(e[0]); }
        int total = 0, bestMax = 0;
        for (int v : nums) { total += v; bestMax = max(bestMax, v); }
        for (int parts = n; parts >= 1; parts--) {
            if (total % parts != 0) continue;
            target = total / parts;
            if (target < bestMax) continue;
            ok = true;
            if (dfs(0, -1) == 0 && ok) return parts - 1;
        }
        return 0;
    }
};

Explanation

If the final components all have the same sum target, then the number of components is total / target, where total is the sum of all values. So target must be a divisor of total, and it must be at least the largest single value (no component can be smaller than its biggest node).

To maximize deleted edges we want as many components as possible, which means the smallest valid target. We try component counts from large to small (equivalently target from small to large) and return on the first one that works.

To test a candidate target, root the tree and DFS. Each node returns the running sum of its still-attached subtree. When a subtree's sum reaches exactly target, we "cut it off" by returning 0, so its parent starts a fresh component. If any subtree sum ever exceeds target, the split is impossible.

A target is achievable exactly when the root's leftover sum is 0 (everything packed neatly into full components) and no subtree overflowed. The number of cuts is then components - 1.

Example: nums = [6,2,2,2,6] has total 18 and max 6. Target 6 yields 3 components {0}, {1,2,3}, {4}, so we delete 3 - 1 = 2 edges.

Time: O(n · d(total)) Space: O(n)