Codility上的问题 12 Lambda 2011

给定一棵树，节点数N在[1..10^5]，树是一维数组给的 T[i] != i表示有一条T[i]到i的边，求定义一个距离是某节点到所有节点距离的和的平均，求这个平均距离最小的节点，如果有相同的输出标号较小的节点。

要求的时间复杂度 O(N),空间复杂度O(N)。

分析： 首先，因为每个节点到其他（N-1)个节点都有路，所以平均距离的话分母是(N - 1)，所以我们只算总距离就好了。另外就是总距离可能非常大，比如一条链……所有节点到根的总距离会超int...

然后如何求这个总距离？ 显然dfs...

先求以一个节点为根节点的子树种的所有节点到它的距离之和。 设s[i]表示以i为根的子树种的节点个数，d[i]表示以i为根的子树的所有节点到i的距离之和。

假设我们要求d[i]，考虑i的一个孩子j, d[j], s[j]已经求好了，那么有s[i] += s[j].距离的话 d[i] += d[j] + s[j]。 这是因为i是j的父亲,j的所有后裔（包括j）到i的距离就是到j的距离+1。

但是我们目前求得的不是所要的结果。

我们现在更新d，让d变为所有节点到它的距离和。

首先对于根节点，d值不用更新。那么我们假设对于节点 j，假设它的父亲i，d[i]已经是所有节点到它的距离和了，那么d[j]怎么办？

我们从整个树中，把以j为根的子树去掉，剩下的节点到i的距离之和为d[i] - (d[j] + s[j]),注意这时d[i],d[j]的含义已经不一样了。那么这些节点到j的距离都要经过其父亲i，也就是都要+1，所以它们到j的距离之和是d[i] - (d[j] + s[j]) + ( n - s[j])，再加上以j自身的子树本身的距离和，所以所求的

d[j] = d[i] - (d[j] + s[j]) + (n - s[j]) + d[j]  = d[i] - s[j] + n - s[j]

这样再dfs一次，所得到的d就是想要的了……

代码：

// you can also use includes, for example:
// #include <algorithm>

#include <vector>

void dfs1(int at, const vector<vector<int> > &e, vector<long long> &d,vector<int> &s) {
int i;
    s[at] = 1;
    d[at] = 0;
    for (i = 0; i < e[at].size(); ++i) {
       dfs1(e[at][i], e, d, s);
       d[at] += d[e[at][i]] + s[e[at][i]];
       s[at] += s[e[at][i]];
    }
}

void dfs2(int at,const vector<vector<int> > &e, vector<long long> &d, vector<int> &s) {
int i;
    for (i = 0; i < e[at].size(); ++i) {
        d[e[at][i]] = d[at] - s[e[at][i]] + s.size() - s[e[at][i]];
        dfs2(e[at][i], e, d, s);
    }
}
       

int solution(const vector<int> &T) {
    // write your code here...
    int n = T.size(),i,r,root;
    vector<vector<int> > e;
    e.resize(n);
    vector<long long> d;
    d.resize(n);
    vector<int> s;
    s.resize(n);
    for (i = 0; i < n; ++i) {
        if (T[i] != i) {
            e[T[i]].push_back(i);
        }
        else {
            root = i;
        }
    }
    dfs1(root, e,d,s);
    dfs2(root, e,d,s);
    for (i = 1, r = 0; i < n; ++i) {
        if (d[r] > d[i]) {
            r = i;
        }
    }
    return r;
        
}