LintCode 1637: Tree problem (经典题)

原题如下：
1637. Tree problem
Given a tree of n nodes. The ith node’s father is fa[i-1] and the value of the ith node is val[i-1]. Especially, 1 represents the root, 2 represents the second node and so on. We gurantee that -1 is the father of root(the first node) which means that fa[0] = -1.
The average value of a subtree is the result of the sum of all nodes in the subtree divide by the number of nodes in it.
Find the maximum average value of the subtrees in the given tree, return the number which represents the root of the subtree.

Example
Given fa=[-1,1,1,2,2,2,3,3], representing the father node of each point, and val=[100,120,80,40,50,60,50,70] , representing the value of each node, return 1.

Sample diagram：
-1 ------No.1
/
No.2 ----1 1---------No.3
/ | \ /
2 2 2 3 3
No.1 node is (100+120+80+40+50+60+50+70) / 8 = 71.25
No.2 node are (120 + 40 + 50 + 60) / 4 = 67.5
No.3 node are (80+50+70) / 3 = 66.6667
So return 1.
Notice
the number of nodes do not exceed 100000
If there are more than one subtree meeting the requirement, return the minimum number.

我的解法就是用map，先遍历一遍两个vector(fa和val)，记下每个节点的sum和count(用mp map)，然后再反向遍历map，将每个节点的sum和count加到其父节点上()。然后又遍历一遍map，记下max average。

#include <limits>

struct Node {
    int sum;
    int count;
};

class Solution {
public:
    /**
     * @param fa: the father
     * @param val: the val
     * @return: the biggest node
     */
    int treeProblem(vector<int> &fa, vector<int> &val) {
        int len = fa.size();
        if (len == 0) return 0;

        map<int, Node> mp;  //id, sum_count
        
        mp[0].sum = val[0];
        mp[0].count = 1;
        
        for (int i = 1; i < len; ++i) {
            mp[i].sum = val[i];
            mp[i].count = 1;
        }
        
        //add the sum and count to the father node
        for (int i = len - 1; i > 0; --i) {  //do not consider fa[0]=-1
            mp[fa[i] - 1].sum += mp[i].sum;
            mp[fa[i] - 1].count += mp[i].count;
        }
 
        int maxId = -1;
        double maxAve = -1 * numeric_limits<double>::max();  //from <limits>

        for (int i = 0; i < len; ++i) {
            double ave = mp[i].sum / mp[i].count;
            if (ave > maxAve) {
                maxId = i;
                maxAve = ave;
            }
        }
        return maxId + 1;
    }
};