백준 2668 - 숫자고르기

문제

https://www.acmicpc.net/problem/2668

리뷰

DFS와 문제 조건을 적절히 그래프로 형성해주면 쉽게 풀이할 수 있는
문제였다.

처음에 이 문제를 접하였을 땐 $N$의 수중에 $k$개를 선택하는 DFS를
이용한 조합을 활용하는 접근으로 문제를 폴이했었다.
하지만, 채택된 0행의 값과 1행의 값을 중복 없이 비교하기 위해 Set을
활용하고 답안을 출력하기 위해 그것을 List로 변환하는 구성이다 보니
메모리 초과가 발생하여 다른 풀이로 방향을 전환했다.

문제 조건을 유심히 보면 0행의 값과 1행의 값을 정점으로 간주하고 같은
열의 두 정점을 간선 관계로 표현하였을때 1->3->1 와 같이 사이클 형태를
띄게 되면 채택 조건을 충족하는 것을 확인할 수 있다.
따라서, visited 배열을 초기화하며 각 정점에서 dfs 를 수행하고
시작 정점으로 돌아오는 사이클을 발견할 시 정답 List에 누적해주는
방식으로 코드를 구성했다.

로직의 시간복잡도는 dfs가 반복문 안에서 실행되고 있으므로 $O(N^2)$으로
수렴하고 $N$이 100인 최악의 경우에도 무난히 제한 조건을 통과한다.

코드

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;

import static java.lang.Integer.*;

public class Main {
    static int N;
    static boolean[] visited;
    static List<List<Integer>> graph = new ArrayList<>();
    static List<Integer> answer = new ArrayList<>();

    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        N = parseInt(in.nextLine());
        visited = new boolean[N + 1];

        for (int i = 0; i <= N; i++)
            graph.add(new ArrayList<>());

        for (int i = 1; i <= N; i++) {
            int value = parseInt(in.nextLine());
            graph.get(i).add(value);
        }

        for (int start = 1; start <= N; start++) {
            dfs(start, 0, start);
            Arrays.fill(visited, false);
        }

        System.out.println(answer.size());
        answer.forEach(System.out::println);
        in.close();
    }

    static void dfs(int current, int depth, int start) {
        visited[current] = true;

        for (Integer next : graph.get(current)) {
            if (next == start) {
                answer.add(start);
                return;
            }

            if (!visited[next]) {
                dfs(next, depth + 1, start);
            }
        }
    }
}

결과