[Algorithm] 백준. 1949 우수마을

www.acmicpc.net/problem/1949

1949번: 우수 마을

 

트리 & 다이나믹 프로그래밍

껄끄러운 두 가지 방식이 합쳐진 문제.

양방향 그래프이므로, DFS방식으로 탐색하면서 visit 체크를 해서 이후에 되돌아 오는 일이 없도록 했다.

리프노드까지 도착한 후에 올라오면서 선택 시 최대 인원, 비선택 시 최대 인원을 저장하는 방식으로 구현하였다.

조건을 만족하려면, 현재 마을인원을 추가하는 경우엔 이웃한 마을의 인원은 추가되면 안되기 때문에, 이웃한 마을들의 비선택시 최대 인원을 누적 인원 + 현재 마을 인원이 최대가 된다. 현재마을을 우수마을로 선택하지 않는 경우엔 이웃한 마을이 선택되던 말던 상관없다. 따라서 선택/비선택 중 큰 값을 누적인원이 비선택시 최대인원이 된다.

코드는 다음과 같다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.LinkedList;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Main {
    static BufferedReader br;
    static StringTokenizer st;
    static int N;
    static boolean [] visit;
    static int [] feed;
    static int [][] cache;
    static LinkedList<Integer>[] list;
    static int max;

    public static void main(String[] args)throws IOException {
        init();
        solve(1);
        max = Math.max(cache[0][1],cache[1][1]);
        System.out.println(max);
    }

    private static void solve(int town) {
        visit[town] = true;
        int rSum = 0;
        int bSum = 0;
        for(Integer each : list[town]){
            if(visit[each])
                continue;
            visit[each]=true;
            solve(each);
            rSum+= Math.max(cache[0][each],cache[1][each]);
            bSum+= cache[0][each];
        }
        cache[0][town] = rSum;
        cache[1][town] = feed[town]+bSum;

    }

    private static void init() throws IOException{
        br = new BufferedReader(new InputStreamReader(System.in));
        st = new StringTokenizer(br.readLine());
        N = StoI();
        visit = new boolean[N + 1];
        feed = new int[N + 1];
        cache = new int[2][N + 1]; // 0 비선택 , 1 선택
        max = 0;
        list = new LinkedList[N + 1];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i < N + 1; i++) {
            list[i] = new LinkedList<>();
            feed[i] = StoI();
        }
        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(br.readLine());
            int from = StoI();
            int to = StoI();
            list[from].add(to);
            list[to].add(from);
        }
    }

    private static Integer StoI() {
        return Integer.parseInt(st.nextToken());
    }
}

처음에는 HashMap이나 TreeSet같은 자료구조로 해결하려고 했으나, 뾰족한 수가 떠오르지 않았고, 선택/비선택의 2차원 캐시 배열을 이용해 간단하게 구현할 수 있었다.

문제 조건이 트리로 주어진 경우엔 DP도 한 번쯤 떠올려 봐야겠다.