Codeforces Round #635 (Div. 2) C. Linova and Kingdom

문제

왕궁에는 n개의 도시가 있고, n-1개의 양방향 도로를 도시에 연결한다. 도시는 1번부터 n번까지 숫자가 붙여져있다. 1번 도시가 수도이며, 왕국은 트리 구조이다.

여왕은 정확히 k의 도시에 산업을 발전시키며 다른 도시들은 관광산업을 발전시킨다. 회의는 1년에 한번 수도에서 진행되는데, 각 산업 도시에서는 대표들을 수도로 보낸다. 대표들은 최단 경로를 따라서 수도에 이동한다.

관광도시를 여행하는 것은 즐겁고, 각 대표의 행복도는 수도에 이동하는 동안 방문한 관광도시의 개수이다. 여왕은 k의 산업도시를 잘 배치해 대표들의 행복도의 합을 최대로 하고 싶다. 

풀이

산업도시가 수도와 얼마나 떨어져 있는지, 즉 depth 값이 각 산업도시의 행복도이다. DFS로 모든 노드를 방문하여, Depth의 내림차순으로 정렬하여 순차적으로 산업도시를 배치하면 될 것 같지만 한가지 생각해주어야 할 점이 있다.

어떤 노드에 산업도시가 생기게 되면, child node의 산업도시들의 행복도가 전부 1씩 감소하게 된다. 따라서 depth 값에서 child node의 개수를 빼준 값이 해당 노드의 실제 행복도 증가가 된다. 

이 수치를 정렬하여 큰 값부터 산업도시를 배치하면 된다.

#include <iostream>
#include <string>
#include <vector>
#include <stack>
#include <algorithm>

using namespace std;

vector<int> roads[200010];
bool visited[200010];
int depthList[200010];

int dfs(int node, int depth)
{
    visited[node] = true;
    int numChild = 0;
    for (int i = 0; i < roads[node].size(); i++)
    {
        int ele = roads[node][i];
        if (visited[ele] == false)
        {
            numChild += dfs(ele, depth + 1) + 1;

        }
    }
    //child node의 개수만큼 손해를 본다.
    depthList[node] = depth - numChild;
    return numChild;
}
bool desc(int a, int b){
    return a>b;
}
int main()
{
    int n, k;
    cin >> n >> k;

    for (int i = 0; i < 200010; i++)
    {
        visited[i] = false;
    }
    for (int i = 0; i < n - 1; i++)
    {
        int c1, c2;
        cin >> c1 >> c2;
        roads[c1].push_back(c2);
        roads[c2].push_back(c1);
    }
    dfs(1, 0);
    sort(depthList + 1, depthList + n + 1, desc);
    long long result = 0;
    for(int i = 1; i <=k; i ++){
        result += depthList[i];
    }
    cout << result << endl;
    return 0;
}