图-1400. 图的费马点

Posted 2021-03-02 hyserendipity

2020-04-06 18:45:31

问题描述：

有一个无向无环连通图，每条边通过两个顶点x[i],y[i]来描述，每条边的长度通过d[i]来描述。
求这样的一个点p，使得其他点到p的距离和最小，如果有多个这样的点p，返回编号最小的。

样例

样例 1:

给出 x = `[1]`, y = `[2]`, d = `[3]`, 返回 `1`。
输入:
[1]
[2]
[3]
输出:
1

解释：
其他点到 1 的距离和为 3，其他点到 2 的距离和为 3，1 的编号较小。

样例 2:

给出 x = `[1,2,2]`, y = `[2,3,4]`, d = `[1,1,1]`, 返回 `2`。
输入:
[1,2,2]
[2,3,4]
[1,1,1]
输出:
2

解释：
其他点到 1 的距离和为 5，其他点到 2 的距离和为 3，其他点到 3 的距离和为 5，其他点到 4 的距离和为 5。

注意事项

• 2 <= n, d[i] <= 10^5
• 1 <= x[i], y[i] <= n

问题求解：

和树的路径和类似，只是这里边有权值。

时间复杂度：O(n)

    Map<Integer, Set<int[]>> graph = new HashMap<>();
    Map<Integer, Integer> cnt = new HashMap<>();
    Map<Integer, Long> sum = new HashMap<>();
    List<long[]> record = new ArrayList<>();
    public int getFermatPoint(int[] x, int[] y, int[] d) {
        for (int i = 0; i < x.length; i++) {
            int from = x[i];
            int to = y[i];
            int w = d[i];
            if (!graph.containsKey(from)) graph.put(from, new HashSet<>());
            if (!graph.containsKey(to)) graph.put(to, new HashSet<>());
            graph.get(from).add(new int[]{to, w});
            graph.get(to).add(new int[]{from, w});
        }
        dfs1(x[0], -1);
        dfs2(x[0], -1);
        Collections.sort(record, (long[] o1, long[] o2) -> o1[0] == o2[0] ? Long.compare(o1[1], o2[1]) : Long.compare(o1[0], o2[0]));
        return (int)record.get(0)[1];
    }
    
    private void dfs2(int root, int parent) {
        record.add(new long[]{sum.get(root), root});
        for (int[] next : graph.get(root)) {
            if (next[0] == parent) continue;
            long w = (long)next[1];
            long curr = sum.get(root) + (long)(cnt.size() - cnt.get(next[0]) * 2) * w;
            sum.put(next[0], curr);
            dfs2(next[0], root);
        }
    }
    
    private void dfs1(int root, int parent) {
        int curr_cnt = 1;
        long curr_sum = 0;
        for (int[] next : graph.get(root)) {
            if (next[0] == parent) continue;
            dfs1(next[0], root);
            curr_cnt += cnt.get(next[0]);
            curr_sum += sum.get(next[0]) + (long)next[1] * (long)cnt.get(next[0]);
        }
        cnt.put(root, curr_cnt);
        sum.put(root, curr_sum);
    }