[sof] 소프티어 <21년 재직자 대회 본선> 거리 합 구하기 (Java) - DFS, 트리

[문제]

https://softeer.ai/practice/info.do?idx=1&eid=635&sw_prbl_sbms_sn=171739 

Softeer

 

[풀이]

DFS를 이용해서 풀 수 있다.

상상도 못한 접근법이라 소프티어가 제공하는 풀이 동영상을 보고서야 풀 수 있었다..참고하면 좋을 듯.

 

*풀이 동영상

Softeer

 

 

먼저 subtree를 이용해야 한다. subtree의 개수는 자기자신을 포함한 자식노드의 개수가 된다. 즉 위 그림에서 2, 5, 6, 7번과 같은 리프노드는 자신밖에 없으므로 subtree의 개수가 1이다. 3번 노드는 3, 5, 6을 포함하므로 3개, 4번 노드는 4, 7을 포함하므로 2개, 1번 노드는 모든 노드를 subtree로 가지므로 7개이다.

 

dfs를 이용해서 subtree의 개수를 구할 수 있고, subtree의 개수를 구하면서 subtree로 연결되어있는 cost들의 합도 같이 구한다. 예를 들어 3번 노드의 subtree sum은 3-5를 잇는 4와 3-6을 잇는 1을 합친 5이다. 1번 노드는 전체 노드를 subtree로 가지고 있으므로 모든 노드와의 거리 cost들을 합한 38이 될 것이다. (5+2+8+6+3+14)

 

이제 dfs를 통해 구한 subtree의 개수와 거리 sum으로 정답을 구할 수 있다. 3번 노드를 예로 들어 보자.

3번 노드의 subtree 개수는 3개이고, sum은 5이다. 부모노드인 1번 노드를 기준으로 전체 노드와의 거리 합을 변경할 수 있다. 3번노드의 subtree에 해당하는 3, 5, 6의 경우 1번 노드에서의 거리보다 cost인 2만큼 줄어든다. 예를 들어 1-5 의 거리는 2+4=6이지만, 3-5의 거리는 2가 줄어든 4이다.

반대로 subtree에 해당하지 않는 1, 2, 4, 7 노드의 경우 1번 노드에서의 거리보다 cost인 2만큼 더 추가된다. 예를 들어, 1-7의 거리는 8+6=14지만 3-7의 거리는 2+8+6=16이 된다.

이를 일반화하면 다음과 같이 나타낼 수 있다.

 

sum[child] = sum[curr] + cost*(n-subtree[child]) - cost*(subtree[child])

실제 코드에서는 위 코드를 정리해서 sum[child] = sum[curr] + cost*(n-2*subtree[child]) 로 나타냈다.

이제 dfs를 돌려서 1번 노드를 기준으로 거리합을 변경해주면 시간초과 없이 문제를 해결할 수 있다!

 

참고로 트리 구성 시 양방향으로 연결해주었기 때문에 dfs 호출 시 curr, parent 노드를 인자로 전달하고, parent!=child일 경우에만 dfs를 호출하여 부모노드에서 리프노드 방향으로만 dfs가 호출되도록 해야 제대로 종료될 수 있다.

 

[코드]

 

import java.util.*;
import java.io.*;


public class Main
{
    static int n;
    static List<Node>[] nodes;
    static long[] subtree;
    static long[] sum;
    static long[] answer; 
    public static class Node{
        int u, cost;
        Node(int u, int cost){
            this.u = u;
            this.cost = cost;
        }
    }  
    public static void main(String args[]) throws Exception
    {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        n = Integer.parseInt(br.readLine());

        nodes = new ArrayList[n+1];
        for(int i=0; i<=n; i++) 
            nodes[i] = new ArrayList<>();
        
        for(int i=1; i<=n-1; i++){
           String[] str = br.readLine().split(" ");
           nodes[Integer.parseInt(str[0])].add(new Node(Integer.parseInt(str[1]),Integer.parseInt(str[2])));
           nodes[Integer.parseInt(str[1])].add(new Node(Integer.parseInt(str[0]), Integer.parseInt(str[2])));
        }

        subtree = new long[n+1];
        sum = new long[n+1];

        dfs1(1, 0); 

        //for(int i=1; i<=n; i++) System.out.println(sum[i]);
        
        dfs2(1, 0);
        for(int i=1; i<=n; i++) System.out.println(sum[i]);
        
    }

    public static void dfs1(int curr, int parent) {
		subtree[curr] = 1;

        for(Node node:nodes[curr]){
            int child = node.u;
            int cost = node.cost;
            if(child!=parent){
                dfs1(child, curr);
                sum[curr] += sum[child] + subtree[child]*cost;
                subtree[curr] += subtree[child]; //자신 포함 서브트리 개수 구하기
            }
        }
	}

    public static void dfs2(int curr, int parent){
        for(Node node:nodes[curr]){
            int child = node.u;
            int cost = node.cost;
            if(child!=parent){
                sum[child] = sum[curr] + cost*(n-2*subtree[child]);
                dfs2(child, curr);
            }
        }
    }
}