【leetcode】(每日抑题 834树中距离之和）

int *ans, *sz, *dp;
int *hd, *nx, *e;

void dfs(int u, int f) {
    
    
    sz[u] = 1;
    dp[u] = 0;
    for (int i = hd[u]; i; i = nx[i]) {
    
    
        int v = e[i];
        if (v == f) {
    
    
            continue;
        }
        dfs(v, u);
        dp[u] += dp[v] + sz[v];
        sz[u] += sz[v];
    }
}

void dfs2(int u, int f) {
    
    
    ans[u] = dp[u];
    for (int i = hd[u]; i; i = nx[i]) {
    
    
        int v = e[i];
        if (v == f) {
    
    
            continue;
        }
        int pu = dp[u], pv = dp[v];
        int su = sz[u], sv = sz[v];

        dp[u] -= dp[v] + sz[v];
        sz[u] -= sz[v];
        dp[v] += dp[u] + sz[u];
        sz[v] += sz[u];

        dfs2(v, u);

        dp[u] = pu, dp[v] = pv;
        sz[u] = su, sz[v] = sv;
    }
}

int* sumOfDistancesInTree(int N, int** edges, int edgesSize, int* edgesColSize, int* returnSize) {
    
    
    ans = malloc(sizeof(int) * N);
    sz = malloc(sizeof(int) * N);
    dp = malloc(sizeof(int) * N);
    hd = malloc(sizeof(int) * N);
    nx = malloc(sizeof(int) * (edgesSize * 2 + 1));
    e = malloc(sizeof(int) * (edgesSize * 2 + 1));
    for (int i = 0; i < N; i++) {
    
    
        ans[i] = sz[i] = dp[i] = hd[i] = 0;
    }
    for (int i = 0, num = 0; i < edgesSize; i++) {
    
    
        int u = edges[i][0], v = edges[i][1];
        nx[++num] = hd[u], hd[u] = num, e[num] = v;
        nx[++num] = hd[v], hd[v] = num, e[num] = u;
    }
    dfs(0, -1);
    dfs2(0, -1);
    *returnSize = N;
    return ans;
}