Question

Google | OA | Remove maximum edges | Set - 25

Answer 1

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

Answer 2

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