16437번 : 양 구출 작전 [JAVA]

문제 링크

성공 여부(걸린 시간): 성공 (59분 40초)

아이디어

이 문제는 트리 구조 위에서 양과 늑대의 개체 수를 누적하면서 최상위 노드(1번)까지 올려보는 방식으로 해결할 수 있다.
내가 정말 어려워 하던 return 이 있는 DFS였는데 사실 재귀를 차분히 생각해서 코드를 짜면 크게 어렵지 않았다.

• 입력으로 주어진 관계는 부모–자식 형태이므로, 자연스럽게 루트(1번 노드)를 기준으로 한 트리 구조를 형성한다.
• 각 노드에는 양 또는 늑대가 주어지고, 양은 양수, 늑대는 음수 값으로 처리한다.
• 리프 노드(더 이상 자식이 없는 노드)에서는 단순히 max(양의 수, 0)을 반환한다. (늑대가 더 많으면 결국 양은 모두 잡아먹히기 때문)
• DFS를 통해 자식 노드부터 계산한 값을 부모 노드로 누적시킨다.
  • 자식 서브트리에서 모아온 양의 수를 부모 노드의 값과 합산한다.
  • 만약 결과가 음수라면(=늑대가 더 많아 양이 전부 잡아먹힌 경우), 해당 서브트리에서는 0만 반환한다.
• 결국, 루트 노드(1번)까지 올라온 최종 합이 섬 전체에서 살아남을 수 있는 양의 수가 된다.

즉, DFS로 트리를 순회하면서, "현재 노드 + 자식 노드의 양들"을 합산하되, 음수가 되면 0으로 잘라주는 방식이다.
이렇게 하면 중복 계산 없이 효율적으로 전체 결과를 구할 수 있다.

---

예제 동작 설명

입력:

7
S 100 1
S 100 1
W 100 1
S 1000 2
W 1000 2
S 900 6

1. 트리 구조

• 1번 노드(루트)
  • 2번 노드 (양 100)
    • 5번 노드 (양 1000)
    • 6번 노드 (늑대 1000)
      • 7번 노드 (양 900)
  • 3번 노드 (양 100)
  • 4번 노드 (늑대 100)

2. DFS 처리 과정

• 리프 5번 노드: 양 1000 → 그대로 1000 반환
• 리프 7번 노드: 양 900 → 그대로 900 반환
• 6번 노드 (늑대 1000): 자기 값(-1000) + 자식(900) = -100 → 음수이므로 0 반환
• 2번 노드 (양 100): 자기 값(100) + 자식(1000 + 0) = 1100 반환
• 3번 노드 (양 100): 그대로 100 반환
• 4번 노드 (늑대 100): 음수이므로 0 반환
• 1번 노드 (루트): 자기 값(0) + 자식(1100 + 100 + 0) = 1200 반환

최종 결과

1200

즉, 이 섬에서는 최종적으로 1200마리의 양이 살아남는다.

풀이 코드

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

public class Main {
    static ArrayList<ArrayList<Integer>> graph;
    static int N;
    static int [] amounts;
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());
        amounts = new int[N+1];
        graph = new ArrayList<>();
        for(int i=0; i<=N; i++){
            graph.add(new ArrayList<>());
        }
        for(int i=2; i<=N; i++){
            StringTokenizer st = new StringTokenizer(br.readLine());
            boolean isSheep = st.nextToken().equals("S");
            int amount = Integer.parseInt(st.nextToken());
            int parent = Integer.parseInt(st.nextToken());
            amount *= isSheep?1:-1;// 양이면 양수 늑대면 음수처리
            graph.get(parent).add(i);
            amounts[i] = amount;
        }
        System.out.print(DFS(1));
    }
    static long DFS(int node){
        if(graph.get(node).isEmpty()){//리프노드라면 늑대일 경우엔 0리턴
            return Math.max(amounts[node], 0);
        }
        long child= amounts[node];
        for(int x:graph.get(node)){
            child += DFS(x);
        }
        return Math.max(0, child);// 음수면(늑대가 더 많다.) 0리턴
    }
}