AOJ - 104 [java] - DFS, BruteForce

 

예제를 보면:

 

 

노드 5개, 색깔: 0 1 1 1 0
간선: 1-2, 2-3, 4-2, 5-4

그림으로 그리면:

 

 

1 - 2 - 3
    |
    4
    |
    5

색깔: 1(빨) - 2(파) - 3(파) - 4(파) - 5(빨)

---

찾아야 하는 것

트리 위의 경로 중에서 딱 4개 노드 A-B-C-D 를 골랐을 때:

조건 1: A→B→C→D가 트리 위의 연속된 경로여야 함

조건 2: 색깔이 빨-파-파-빨 또는 파-빨-빨-파 여야 함

조건 3: (A,B,C,D)랑 (D,C,B,A)는 같은 걸로 봄 (방향만 반대)

---

예제에서 직접 찾기

경로 1-2-4-5를 보면:

 

 

1(빨) - 2(파) - 4(파) - 5(빨)

→ 빨-파-파-빨 ✅ 조건 만족!

다른 경로 1-2-3은 노드가 3개라 탈락, 1-2-4도 3개라 탈락.

그래서 정답이 1

 

 

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

public class Main {

    static int N;
    static int[] color;
    static List<Integer>[] adj;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine().trim());

        color = new int[N + 1];
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) color[i] = Integer.parseInt(st.nextToken());

        adj = new ArrayList[N + 1];
        for (int i = 1; i <= N; i++) adj[i] = new ArrayList<>();
        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());
            adj[u].add(v);
            adj[v].add(u);
        }

        long answer = 0;

        for (int a = 1; a <= N; a++) {
            // a에서 DFS → 모든 노드의 부모 기록
            int[] parent = dfs(a);

            for (int d = a + 1; d <= N; d++) {
                if (color[a] != color[d]) continue;

                // a → d 경로 복원
                List<Integer> path = getPath(parent, a, d);

                if (path.size() < 4) continue;

                int b = path.get(1);
                int c = path.get(path.size() - 2);

                if (color[b] == color[c] && color[b] != color[a]) {
                    answer++;
                }
            }
        }

        System.out.println(answer);
    }

    // a를 시작점으로 DFS, 각 노드의 부모를 기록한 배열 반환
    static int[] dfs(int start) {
        int[] parent = new int[N + 1];
        Arrays.fill(parent, -1);

        boolean[] visited = new boolean[N + 1];
        Stack<Integer> stack = new Stack<>();

        stack.push(start);
        visited[start] = true;

        while (!stack.isEmpty()) {
            int cur = stack.pop();
            for (int next : adj[cur]) {
                if (!visited[next]) {
                    visited[next] = true;
                    parent[next] = cur;  // next의 부모 = cur
                    stack.push(next);
                }
            }
        }

        return parent;
    }

    // parent 배열을 이용해 start → end 경로 복원
    static List<Integer> getPath(int[] parent, int start, int end) {
        List<Integer> path = new ArrayList<>();

        // end에서 parent 따라가면 start까지 거슬러 올라감
        int cur = end;
        while (cur != -1) {
            path.add(cur);
            cur = parent[cur];
        }

        // [end, ..., start] → [start, ..., end] 로 뒤집기
        Collections.reverse(path);

        return path;
    }
}

---

흐름 요약

main()
 ├─ 입력 받기
 ├─ for a (모든 시작 노드)
 │    ├─ dfs(a) → parent[] 반환
 │    │    └─ 스택으로 a에서 출발해 모든 노드 방문
 │    │       방문할 때마다 parent[next] = cur 기록
 │    │
 │    └─ for d (a보다 큰 모든 끝 노드)
 │         ├─ getPath(parent, a, d) → [a, ..., d] 반환
 │         │    └─ d에서 parent 따라 거슬러 올라간 뒤 reverse
 │         │
 │         └─ path[1] = b, path[끝-1] = c 꺼내서 색깔 확인

핵심 두 가지만 기억하면 됩니다.

• dfs() → parent 배열 채우기 (누가 누구의 부모인지)
• getPath() → parent 역추적으로 경로 복원