소프티어 선물 나눠주기 (EASY)

문제 링크

이 문제에서 각 노드 별 가중치는 무시해도 된다. 단 각 노드의 가중치가 0이 아닌 점은 표시를 해야합니다.

가중치가 0이 아닌 노드들을 사이의 최단 탐색 경로의 거리를 구하는 문제입니다.

Tree에는 다음과 같은 속성이 있습니다.

어떤 노드를 Root로 잡아도 tree형태가 된다.

이 속성을 사용해 가중치가 있는 아무 노드를 Root로 설정합니다.

그 다음 DFS로 나머지 가중치가 있는 정점까지의 최소값을 탐색합니다.

이 거리 * 2를 해주면 최단 탐색 경로 거리가 됩니다.

	public static int dfsMark(int cur) {
        visited[cur] = true;
        boolean foundRequired = isRequiredTown[cur];
        int count = 0;
        
        for (int next : adjList[cur]) {
            if (visited[next]) continue;
            int subCount = dfsMark(next);
            if (subCount != 0) {
                foundRequired = true;
                count += subCount;
                // 자식 중 하나라도 필수 마을이 있으면 현재 노드도 필수로 포함
                isRequiredTown[cur] = true;
            }
        }
        return foundRequired ? count + 1 : 0;
    }

쿼리 처리는 노드와 트리 사이의 거리를 구해주면 됩니다.

    public static int dfsFindDistance(int cur, int dist) {
        visited[cur] = true;
        if (isRequiredTown[cur]) {
            return dist;
        }
        
        for (int next : adjList[cur]) {
            if (visited[next]) continue;
            int candidate = dfsFindDistance(next, dist + 1);
            if (candidate != 0) {
                // 경로 상의 노드들을 스테이너 트리에 포함시키기
                isRequiredTown[cur] = true;
                return candidate;
            }
        }
        return 0;
    }

전체 코드

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

public class Main {

    static int n;
    static int[] scores;
    static ArrayList<Integer>[] adjList;
    // 필수 방문해야 하는 마을 여부 (선물이 있거나 배달해야 하는 마을)
    static boolean[] isRequiredTown; 
    static boolean[] visited;

    public static int dfsMark(int cur) {
        visited[cur] = true;
        boolean foundRequired = isRequiredTown[cur];
        int count = 0;
        
        for (int next : adjList[cur]) {
            if (visited[next]) continue;
            int subCount = dfsMark(next);
            if (subCount != 0) {
                foundRequired = true;
                count += subCount;
                // 자식 중 하나라도 필수 마을이 있으면 현재 노드도 필수로 포함
                isRequiredTown[cur] = true;
            }
        }
        return foundRequired ? count + 1 : 0;
    }

    public static int dfsFindDistance(int cur, int dist) {
        visited[cur] = true;
        if (isRequiredTown[cur]) {
            return dist;
        }
        
        for (int next : adjList[cur]) {
            if (visited[next]) continue;
            int candidate = dfsFindDistance(next, dist + 1);
            if (candidate != 0) {
                // 경로 상의 노드들을 스테이너 트리에 포함시키기
                isRequiredTown[cur] = true;
                return candidate;
            }
        }
        return 0;
    }

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        n = Integer.parseInt(br.readLine());
        
        scores = new int[n + 1];
        adjList = new ArrayList[n + 1];
        isRequiredTown = new boolean[n + 1];
        visited = new boolean[n + 1];
        
        for (int i = 1; i <= n; i++) {
            adjList[i] = new ArrayList<>();
        }
        
        StringTokenizer st = new StringTokenizer(br.readLine());
        int startingPoint = 0;
        for (int i = 1; i <= n; i++) {
            scores[i] = Integer.parseInt(st.nextToken());
            if (scores[i] != 0) {
                isRequiredTown[i] = true;
                startingPoint = i;
            }
        }
        
        // 트리 그리기
        for (int i = 0; i < n - 1; i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());
            adjList[u].add(v);
            adjList[v].add(u);
        }
        
        // 초기 스테이너 트리 구성
        Arrays.fill(visited, false);
        int initialCount = dfsMark(startingPoint);
        int totalDistance = (initialCount - 1) * 2;
        System.out.println(totalDistance);
        
        // 쿼리 처리
        int q = Integer.parseInt(br.readLine());
        for (int i = 0; i < q; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            
            // 정점 a가 스테이너 트리에 포함되어 있지 않으면 추가
            if (!isRequiredTown[a]) {
                visited = new boolean[n + 1];
                totalDistance += dfsFindDistance(a, 0) * 2;
                isRequiredTown[a] = true;
            }
            
            // 정점 b가 스테이너 트리에 포함되어 있지 않으면 추가
            if (!isRequiredTown[b]) {
                visited = new boolean[n + 1];
                totalDistance += dfsFindDistance(b, 0) * 2;
                isRequiredTown[b] = true;
            }
            System.out.println(totalDistance);
        }
    }
}