[JAVA] 9934 - 완전 이진 트리 🩶1 : 트리 탐색은 index 기반으로 탐색하자.

들어가며

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

문제의 입력값으로 트리의 높이(K)와 중위순회(InOrder)로 탐색한 노드가 순서대로 주어진다. 이를 레벨별로 다시 구성하여 출력하면 된다. 문제의 요구사항은 단순했는데, 너무 무식하게 풀어서 내 풀이와 다른분들의 풀이를 비교하면서 글을 적어본다.

본문으로

처음에는 InOrder로 순회한 것을 어떻게 레벨별로 다시 구성하지? 라는 생각과 함께 아이디어가 떠오르지 않아서 그려보았다.

역시 직접 손으로 그려보니 바로 아이디어가 보였다. 리스트를 반으로 잘라서 그 절반에 있는 값이 그 레벨에 속하는 값이구나..! 생각이 들었다. 그래서 바로 리스트를 반으로 잘라가며 넘겨야겠다. 라는 생각에 아래와 리스트를 복사해가며 로직을 구성했다.

public class Main {

    static int K;
    static List<Integer>[] answer;

    static void go(int idx, List<Integer> beforeList, List<Integer> afterList) {
        if (beforeList.size() == 1 && afterList.size() == 1) {
            answer[idx].add(beforeList.get(0));
            answer[idx].add(afterList.get(0));
            return;
        }

        answer[idx].add(beforeList.get(beforeList.size() / 2));
        answer[idx].add(afterList.get(afterList.size() / 2));

        List<Integer> aBeforeList = new ArrayList<>();
        List<Integer> aAfterList = new ArrayList<>();
        List<Integer> bBeforeList = new ArrayList<>();
        List<Integer> bAfterList = new ArrayList<>();

        for (int i = 0; i < beforeList.size() / 2; i++) {
            aBeforeList.add(beforeList.get(i));
        }
        for (int i = beforeList.size() / 2 + 1; i < beforeList.size(); i++) {
            aAfterList.add(beforeList.get(i));
        }
        for (int i = 0; i < afterList.size() / 2; i++) {
            bBeforeList.add(afterList.get(i));
        }
        for (int i = afterList.size() / 2 + 1; i < afterList.size(); i++) {
            bAfterList.add(afterList.get(i));
        }

        go(idx + 1, aBeforeList, aAfterList);
        go(idx + 1, bBeforeList, bAfterList);

    }

    public static void main(String[] args) throws IOException {
        BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));
        K = Integer.parseInt(bf.readLine());
        answer = new ArrayList[K];
        for (int i = 0; i < K; i++) {
            answer[i] = new ArrayList<>();
        }

        int size = 2;
        for (int i = 0; i < K - 1; i++) {
            size *= 2;
        }
        size -= 1;

        StringTokenizer st = new StringTokenizer(bf.readLine());
        List<Integer> list = new ArrayList<>();
        for (int i = 0; i < size; i++) {
            list.add(Integer.parseInt(st.nextToken()));
        }

        answer[0].add(list.get(size / 2));
        List<Integer> beforeList = new ArrayList<>();
        List<Integer> afterList = new ArrayList<>();
        for (int i = 0; i < size / 2; i++) {
            beforeList.add(list.get(i));
        }
        for (int i = size / 2 + 1; i < size; i++) {
            afterList.add(list.get(i));
        }

        go(1, beforeList, afterList);

        for (int i = 0; i < K; i++) {
            for (int j = 0; j < answer[i].size(); j++) {
                System.out.print(answer[i].get(j) + " ");
            }
            System.out.println();
        }
    }
}

보면.. 리스트를 슬라이싱하는 법을 몰라서 그냥.. 새로 만들어서 원하는 만큼 담아준 다음에 인자로 넘겨줬다. 다행히도 시간초과는 나지 않고 정답이었다. 하지만 이 문제가 제한시간을 넉넉하게 줘서 다행이지, 매번 List<> 자료구조를 생성하는 것은 비효율적인 것 같았다.

---

그래서 다른 분들 풀이를 참고해보니, 아니나 다를까 나처럼 리스트를 통째로 넘기는 분은 거의 없고 대부분 원본 리스트에서 인덱스를 가지고 구현하셨다... 왜 그생각을 못했을까? 아래는 참고하여 인덱스 기반으로 다시 구현한 로직이다.

public class Main {

    static int K;
    static List<Integer>[] answer;
    static List<Integer> tree;

    static void go(int level, int start, int end) {
        if (level == K || start > end) {
            return;
        }

        int mid = (start + end) / 2;
        answer[level].add(tree.get(mid));

        go(level + 1, start, mid - 1);
        go(level + 1, mid + 1, end);
    }

    public static void main(String[] args) throws IOException {
        BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));
        K = Integer.parseInt(bf.readLine());
        answer = new ArrayList[K];
        for (int i = 0; i < K; i++) {
            answer[i] = new ArrayList<>();
        }

        int size = 2;
        for (int i = 0; i < K - 1; i++) {
            size *= 2;
        }
        size -= 1;

        StringTokenizer st = new StringTokenizer(bf.readLine());
        tree = new ArrayList<>();
        for (int i = 0; i < size; i++) {
            tree.add(Integer.parseInt(st.nextToken()));
        }

        go(0, 0, size - 1);

        for (int i = 0; i < K; i++) {
            for (int j = 0; j < answer[i].size(); j++) {
                System.out.print(answer[i].get(j) + " ");
            }
            System.out.println();
        }
    }
}

이렇게 인덱스 기반으로 한다면 각 원소를 answer배열에 딱 한번만 넣으므로 O(N)으로 시간복잡도를 유지할 수 있다. 원래 내 풀이는.. 레벨마다 리스트를 복사하므로 대략 O(N log N) 정도로 시간복잡도가 불어난 풀이었다.

마무리하며

트리 관련된 문제가 나오면 괜히 어려워 보이는데, 인덱스 기반으로 잘 생각해보자! 트리 문제는 전공자라면 잘 풀어야 하지 않겠는가!!