Trees: Find the Minimum Difference of two floods

A sewer drainage system is structured as a tree. Water enters the system at n nodes numbered from 0 to n-1 and flows through the tree to the root, which has the number 0.

The tree structure is defined by an array parent, where parent[i] = j means that water flows from node i to its direct parent node j. Water exits the system after it flows through the root, so the root has no parent, represented as parent[0] = -1. The value in input[i] denotes the amount of water that enters the sewer system at node i. This excludes water that flows into i from its children.

The total flow through a node is thus the flow that enters at that node, plus the sum of the total flows of all of its children.

Your task is to divide the system into two smaller pieces by cutting one branch so that the total flows of the resulting subtrees are as close as possible.

Example:

parent = [-1, 0, 0, 1, 1, 2] 
input = [1, 2, 2, 1, 1, 1]

Cut the branch between nodes 1 and 0. 

The partition {0, 2, 5) has total flow input[0] + input[2] + input[5] = 1+2+1 = 4. 
The partition {1, 3, 4) has total flow input[1] + input[3] + input[4] = 2 + 1 + 1 = 4. 

The absolute difference between the total flows of the two new sewer systems is 4 -4 = 0. It's not possible for a different division to achieve a smaller difference than 0, so the final answer is 0. 

Solution

We build the map using the array, then we build the tree using a recursive algorithm that adds the sum of children to the root. To find the connection that will give the minimum difference between the floods, we simply go through all the elements in the map and try to remove its substree.

Since the root is the sum of all elements the difference, the difference between that subtree and the rest is:

=flood of everything but subtree(sum in root - sum in element) - substree flood(sum in element)

class Node{
public:
	int val;
    vector<Node*> children;
    Node(int key) {
    	val = key;
    }
};

void findMinDiff(Node* root, int rt, int& m) {
	if(root != NULL) {
    	m = min(m, rt - 2 * root->val);
        for(auto n: root->children)
        	findMinDiff(n, rt, m);
    }
}

int solution(vector<int> parent, vector<int> input) {
	Node* root = buildTree(parent, input);
    flood(root);
    int m = INT_MAX;
    for(auto n: root->children)
    	findMinDiff(n, root->val, m);
	return m;
}

void buildTree(vector<int>& parent, vector<int>& input) {
	vector<Node*> nodes(parent.size(), nullptr);
    int r;
    for(int i = 0; i < parent.size(); i++) {
    	if(nodes[i] == nullptr)
        	buildNodeUp(parent, input, nodes, i, r);
    }
    return nodes[r];
} 

void buildNodeUp(vector<int>& parent, vector<int>& input, vector<Node*> nodes, int i, int& r) {
	nodes[i] = new Node(input[i]);
    if(parent[i] == -1) {
    	r = i;
        return;
    }
    
    if(nodes[parent[i]] == nullptr)
    	buildNodeUp(parent, input, nodes, parent[i], r);
    nodes[parent[i]]->children.push_back(nodes[i]);
}

void flood(Node* root) {
	int res;
    if(root != nullptr) {
    	int sm = 0;
        for(auto n: root->children)
        	sm += flood(n);
        root->val += sm;
    }
}