洛谷题单指南-二叉树-P1364 医院设置

原题链接：https://www.luogu.com.cn/problem/P1364

题意解读：医院的位置使所有居民所走的路程之和为最小，即找到一个节点，该节点到其他所有节点的距离*其他节点的权值即人数之和最小。

解题思路：

看起来是一个二叉树问题，本质上是一个图论问题，有三种方式可以求解：

1、计算每两个节点之间的最短路径（Floyd算法），然后枚举在每一个点建医院，计算所有居民的路程之和，求最小值，复杂度为O(n^3)，这里不做考虑。

2、从每一个节点出发，进行DFS搜索，搜索过程中计算每个点到起始点的路程*人数，求和，对和取最小值。

3、扩展：如果n的范围变成10000，以上O(n^2)以上的算法将都失效，需要借助树的重心概念来求解。

方法一、O(n^2)算法-DFS搜索

要实现从每一个节点出发进行BFS或者DFS，需要存储每个节点的子节点、父节点信息

struct node
{
    int weight; //人口数
    int sub[3]; //sub[0]父节点、sub[1]左子节点、sub[2]右子节点
} tree[N];

下面给出DFS版本的代码，因为个人感觉DFS代码更简洁（但BFS代码没有递归更容易理解）。

100分代码：

#include <bits/stdc++.h>
using namespace std;

const int N = 105;

struct node
{
    int weight; //人口数
    int sub[3]; //sub[0]父节点、sub[1]左子节点、sub[2]右子节点
} tree[N];

bool flag[N];

int n, w, u, v;
int sum, ans = INT_MAX;

//idx是节点号,当前节点距离起始节点路径长度depth,路程和sum
void dfs(int idx, int depth)
{
    flag[idx] = true;
    if(depth > 0) sum += tree[idx].weight * depth;
    for(int i = 0; i < 3; i++)
    {
        if(!flag[tree[idx].sub[i]]) dfs(tree[idx].sub[i], depth + 1);
    }
}

int main()
{
    cin >> n;
    for(int i = 1; i <= n; i++)
    {
        cin >> w >> u >> v;
        tree[i].weight = w;
        tree[i].sub[1]  = u;
        tree[i].sub[2] = v;
        if(u) tree[u].sub[0] = i; //记录u的父节点为i
        if(v) tree[v].sub[0] = i; //记录v的父节点为i
    }

    for(int i = 1; i <= n; i++) 
    {
        sum = 0; //路程之和初始化
        memset(flag, 0, sizeof(flag)); //flag初始化为false
        dfs(i, 0);
        ans = min(ans, sum);
    }
    cout << ans;

    return 0;
}

方法二、O(n)算法-树的重心+树上DP

可以使用树形动态规划来解决这个问题。核心思路是先任选一个节点作为根节点，将树转化为有根树，然后通过两次DFS完成动态规划过程：

第一次 DFS：计算以根节点为医院时，所有居民到医院的距离之和，同时记录每个节点的子树中居民的总数。

第二次 DFS：从根节点开始，通过状态转移，计算将医院设置在每个节点时，所有居民到医院的距离之和。

记录每个节点的子树中居民的总数方法，可以借鉴树的重心概念，请参考：https://www.cnblogs.com/jcwy/p/18750178

由于每个节点是有权值（代表居民人口），在计算子树节点数量时要考虑权值。

设sz[i]表示以节点i为根的子树的节点数，任取节点1为起点开始DFS，不难得到sz[]数组；

设f[i]表示以i为根建立医院，所有居民走的总距离，f[1]值在DFS过程中也可以计算出；

DFS代码为：

//从起点1开始，计算以每个点为根的子树大小sz[u]，同时计算1到所有点的距离之和f[1]
void dfs1(int u, int parent, int depth)
{
    f[1] += w[u] * depth;
    sz[u] = w[u];
    for(auto v : g[u])
    {
        if(v == parent) continue;
        dfs1(v, u, depth + 1);
        sz[u] += sz[v];
    }
}

最终要求f[i]的最大值，已知f[1]，可以通过递推计算其余f[i]值，关键来看f值之间有什么递推关系。

设节点u->v，已知f[u]，也就是以u为根节点的总距离，如何计算f[v]呢？

当从u走到v，所有以v为根的子树节点原本要走到u，现在只用走到v，所有以v为根的子节点距离都减少1，总距离减少了sz[v]；

而所有以v为根的子树以外的节点原本只用走到u，现在却要走到v，这些节点距离都增加1，总距离增加了total - sz[v];

因此f[v] = f[u] - sz[v] + total - sz[v]，total是总节点数（居民人数，等同于sz[1]）。

100分代码：

#include <bits/stdc++.h>
using namespace std;

const int N = 10005;

vector<int> g[N];
int w[N]; //每个节点的人口数
int sz[N]; //sz[i]表示以节点i为根的子树的节点数
int f[N]; //f[i]表示以i为根建立医院，所有居民走的总距离
int n;
int ans = INT_MAX;

//从起点1开始，计算以每个点为根的子树大小sz[u]，同时计算1到所有点的距离之和f[1]
void dfs1(int u, int parent, int depth)
{
    f[1] += w[u] * depth;
    sz[u] = w[u];
    for(auto v : g[u])
    {
        if(v == parent) continue;
        dfs1(v, u, depth + 1);
        sz[u] += sz[v];
    }
}

//通过树形关系递推根据f[u]求f[v]
void dfs2(int u, int parent)
{
    for(auto v : g[u])
    {
        if(v == parent) continue;
        f[v] = f[u] - sz[v] + sz[1] - sz[v];
        dfs2(v, u);
    }
}

int main()
{
    cin >> n;
    int x, v1, v2;
    for(int u = 1; u <= n; u++)
    {
        cin >> w[u] >> v1 >> v2;
        if(v1) 
        {
            g[u].push_back(v1);
            g[v1].push_back(u);
        }
        if(v2) 
        {
            g[u].push_back(v2);
            g[v2].push_back(u);
        }
    }

    dfs1(1, 0, 0);
    dfs2(1, 0);

    for(int i = 1; i <= n; i++) 
        ans = min(ans, f[i]);

    cout << ans;

    return 0;
}