[Java] 백준 15663번: N과 M (9)

[문제 바로 가기] - N과 M (9)

유형 : 백트래킹

💡 아이디어

중복되는 수열이 있으면 안된다고 해서 Set, 그 중에서도 수열이 사전 순으로 증가해야 한다고 해서 TreeSet을 사용하기로 했다.

처음 풀이(1128ms)

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Set;
import java.util.StringTokenizer;
import java.util.TreeSet;

public class Main {
    static int[] nums, select;
    static int n, m;
    static Set<int []> result = new TreeSet<>((a, b) -> {
        for (int i = 0; i < m; i++) {
            if (a[i] != b[i]) return a[i] - b[i];
        }
        return 0;
    });
    static boolean[] visited;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        n = Integer.parseInt(st.nextToken());
        m = Integer.parseInt(st.nextToken());

        nums = new int[n];
        st = new StringTokenizer(br.readLine());

        for (int i = 0; i < n; i++) {
            nums[i] = Integer.parseInt(st.nextToken());
        }

        select = new int[m];
        visited = new boolean[n];

        perm(0);

        for (int[] n : result) {
            for (int i = 0; i < m; i++) {
                System.out.print(n[i] + " ");
            } System.out.println();
        }
    }

    public static void perm(int count) {
        if (count == m) {
            result.add(select.clone());
            return;
        }

        for (int i = 0; i < n; i++) {
            if (!visited[i]) {
                visited[i] = true;
                select[count] = nums[i];
                perm(count + 1);
                visited[i] = false;
            }
        }
    }
}

순열을 사용해서 풀이했고 TreeSet을 사용해서 그런지 실행 시간이 정말 오래 걸렸다.

TreeSet은 레드 블랙 트리로 이루어져 있기 때문에 원소 하나를 넣을 때마다 트리 위치를 찾고, 그 과정에서 Comparator로 계속 비교하고, 균형을 맞추기 위해 회전 연산을 한다.

심지어 int[] 타입이기 때문에 m번을 추가로 비교하기 때문에 O(m)의 시간 복잡도가 추가로 발생한다.

for (int i = 0; i < m; i++) {
    if (a[i] != b[i]) return a[i] - b[i];
}

따라서 Set을 사용하는 대신 백트래킹 연산 과정에서 같은 수가 들어가면 스킵하도록 코드를 수정했다.

이때 nums를 먼저 정렬해주면 이후에 정렬할 필요 없이 사전 순으로 수열을 만들 수 있다.

120ms 풀이

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {
    static int[] nums, select;
    static int n, m;
    static boolean[] visited;
    static StringBuilder sb = new StringBuilder();

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        n = Integer.parseInt(st.nextToken());
        m = Integer.parseInt(st.nextToken());

        nums = new int[n];
        st = new StringTokenizer(br.readLine());

        for (int i = 0; i < n; i++) {
            nums[i] = Integer.parseInt(st.nextToken());
        }

        select = new int[m];
        visited = new boolean[n];

        Arrays.sort(nums);

        perm(0);

        System.out.println(sb);
    }

    public static void perm(int count) {
        if (count == m) {
            for (int i = 0; i < m; i++) {
                sb.append(select[i]).append(" ");
            }

            sb.append("\n");
            return;
        }

        int prev = -1;

        for (int i = 0; i < n; i++) {
            if (!visited[i]) {

                if (prev == nums[i]) continue;

                visited[i] = true;
                select[count] = nums[i];
                prev = nums[i];

                perm(count + 1);

                visited[i] = false;
            }
        }
    }
}