0 like 0 dislike
1,522 views
| 1,522 views

0 like 0 dislike

Problem statement:

Given a tree with N nodes and N-1 edges. Each node has a value associated with it.
You have to remove maximum number of edges such that the value of each resulted connected component is same.
The value of a connected component is the sum of values of it's nodes.

Note: Answer always exists as you can always remove zero edges, so that the whole tree remains connected.

Input Format:

First line contains integer N (number of nodes)
Second line contains N integers, where ith integer is value of ith node.
Next N-1 lines contains 2 integers u, v denoting there is an edge between node u and v of tree.

Constraints:

1 <= N <= 10^5
1 <= a[i] <= 10^3

Output:

The maximum number of edges you can remove from the tree such that it holds the condition given in the problem.

Example:

Input:
5
3 2 1 1 1
1 2
1 3
2 4
2 5

Output:
1

Solution:
(Based on ideas from @nvdrag123 , @anarchaworld and @LC150H100M )

Basic idea:

• Find all factors of sum, sum will be at max 10^8, so any number <= 10^8 will have at max 100 factors
• For all factors, check how many components we can break our tree into
• Start from the smallest factor and get first valid factor which satisfies above condition (small factor => more components => more edges removed)
• To check valid division, recursively calculate number of components for each subtree by pre computing subtree sum as shown in code below
by Expert (111,330 points)
0 like 0 dislike
```vector<vector<int>> graph;
vector<int> nums; // value of nodes
vector<int> subtreeSum; // sum of subtree rooted at current node
vector<int> componentsBelow; // components made in subtree of current node, for given factor
int targetSum, totalComponents;

// recursively calculating subtree sum rooted at each node
void calculateSubtreeSum(int current, int parent) {
subtreeSum[current] = nums[current];
for(int child : graph[current]) {
if(child != parent) { // go only to children subtrees of current node
calculateSubtreeSum(child, current);
subtreeSum[current] += subtreeSum[child];
}
}
}

// count components that can be made in subtree of current node,
// such that all have sum of values = targetSum
void countComponents(int current, int parent) {
// division at current node
if(subtreeSum[current] == targetSum) {
componentsBelow[current] = 1;
totalComponents++;
return;
}

// count all divisions in subtrees
for(int child : graph[current]) {
if(child != parent) {
countComponents(child, current);
componentsBelow[current] += componentsBelow[child];
}
}

// sum of nodes left in subtree after component divisions == targetSum means one more division
if(subtreeSum[current] - (componentsBelow[current] * targetSum) == targetSum) {
componentsBelow[current]++;
totalComponents++;
}
}

int main() {
// input and driver code
int n, totalSum = 0, u, v;
vector<int> factors;

cin >> n;
nums.resize(n+1);
graph.resize(n+1, vector<int>());
subtreeSum.resize(n+1);

for(int i = 1; i <= n; ++i) {
cin >> nums[i];
totalSum += nums[i];
}

for(int i = 1; i < n; ++i) {
cin >> u >> v;
graph[u].push_back(v);
graph[v].push_back(u);
}

// calculating all factors of totalSum
for(int factor = 1; factor <= sqrt(totalSum); ++factor) {
if(totalSum % factor == 0) factors.push_back(factor);
if(factor != (totalSum / factor)) factors.push_back(totalSum / factor);
}
sort(factors.begin(), factors.end());

calculateSubtreeSum(1, 0); // pre computation of subtree sum

// we start from smallest factor, and stop when we get any valid division
// because smallest factor => most components, so it's optimal
for(int factor : factors) {
targetSum = factor;
totalComponents = 0;
componentsBelow.assign(n+1, 0);

countComponents(1, 0); // pre computation of valid components for each subtree
// if this division is valid
if(totalComponents * factor == totalSum) {
// we need to remove x-1 edges to break tree into x components
cout << totalComponents - 1 << '\n';
return 0;
}
}

// if no division was valid, we don't remove any edge
cout << "0\n";
return 0;
}
```
by Expert (111,330 points)