我们的距离

思路

最起初显然可以得到一个时间复杂度为 O(n2)O(n^2) 的暴力做法,思考一下只会会发现两条性质:

  • 当我们固定某个结点为根节点时,对于结点 xx,设 dis[x]dis[x] 为结点 xx 与其子树中所有结点的距离和,sz[x]sz[x] 为其子树的大小,则易得:dis[x]=v(dis[v]+sz[v])dis[x] = \sum_{v}(dis[v]+sz[v])

  • f[x]f[x] 为结点 xx 的权值,若 xxvv 的父节点,则有关系 f[v]=dis[v]+(f[x]dis[v]sz[v])+nsz[v]=f[x]+n2×sz[v]f[v] = dis[v] + (f[x] - dis[v] - sz[v]) + n - sz[v] = f[x] + n - 2 \times sz[v]

结合上述两条性质,我们只需要先做一次 dfs 预处理一下 disdis 数组,之后再用 dfs 结合上述两条性质处理一下每个点的权值即可。 时间复杂度 O(n)O(n)

Code

class Solution {
#define pii pair<int, int>
#define ll long
#define maxn 100010
private:
    int n;
    vector<int> g[maxn];
    ll dis[maxn], sz[maxn], f[maxn];
public:
    /**
     * 返回最终结果
     * @param n int整型 城市数量
     * @param u int整型vector
     * @param v int整型vector
     * @param w int整型vector
     * @return long长整型
     */
 
    void dfs1(int x, int fa = 0) {
        for (auto v: g[x]) {
            if (v == fa) continue;
            dfs1(v, x);
            sz[x] += sz[v];
            dis[x] += sz[v] + dis[v];
        }
    }
    void dfs2(int x, int fa = 0) {
        for (auto v: g[x]) {
            if (v == fa) continue;
            ll siz = n - sz[v];
            f[v] = dis[v] + (f[x]- dis[v] - sz[v]) + siz;
            dfs2(v, x);
        }
    }
    vector<ll> solve(int n, vector<Point>& edge) {
        // write code here
        this->n = n;
        for (int i = 1; i <= n; ++i) {
            g[i].clear();
            f[i] = dis[i] = 0;
            sz[i] = 1;
        }
        for (int a, b, i = 1; i < n; ++i) {
            a = edge[i-1].x, b = edge[i-1].y;
            g[a].push_back(b);
            g[b].push_back(a);
        }
        dfs1(1);
        f[1] = dis[1];
        dfs2(1);
        vector<ll> ans;
        for (int i = 1; i <= n; ++i)
            ans.push_back(f[i]);
        return ans;
 
    }
};

全部评论

相关推荐

不愿透露姓名的神秘牛友
12-03 15:41
已编辑
淘天 算法工程师 31.0k*16.0
点赞 评论 收藏
分享
冷艳的小师弟在看机会:jd测评乱点直接被挂了,哭死~
点赞 评论 收藏
分享
评论
点赞
收藏
分享
牛客网
牛客企业服务