Binary Tree Inorder Traversal (Stack)

LeetCode의 Queue, Stack 튜토리얼에서 Stack과 DFS 섹션의 마지막 문제인 Binary Tree Inorder Traversal다.

Explore - LeetCode

이전까지는 주로 재귀를 이용하여 DFS를 구현했다면 이번에는 명확하게 Stack 자료구조를 활용하여 문제를 풀어보라고 하고 있다. 그럼에도 불구하고 역시 상위권 풀이는 전부 재귀를 활용한 DFS인데 아마 튜토리얼을 진행하면서 푸는 게 아니라 그냥 따로 문제를 푸는 사람들인 것 같다.

 

Inorder는 트리 자료구조에서 트리를 탐색하는 방법 중 하나로 주로 이진트리 탐색을 예시로 볼 수 있다. 크게 4가지(Inorder, Preorder, Postorder, Levelorder) 탐색 방법이 있는데 이는 현재 노드와 자식 노드(2개라고 가정)를 어떻게 탐색하는지 그 순서를 정의한 방식이다. 잘 설명된 문서를 참조하면 다음과 같다.

• Inorder: 왼쪽 자식 노드 탐색, 현재 노드 탐색, 오른쪽 자식 노드 탐색.

• Preorder: 현재 노드 탐색, 왼쪽 자식 노드 탐색, 오른쪽 자식 노드 탐색.

• Postorder: 왼쪽 자식 노드 탐색, 오른쪽 자식 노드 탐색, 현재 노드 탐색.

• Levelorder: 현재 노드와 같은 높이에 있는 모든 노드 탐색 후 자식 노드와 같은 높이에 있는 모든 노드 탐색.

여기서 노드를 탐색한다는 것은 해당 노드를 어떤 로직으로 처리한다는 것을 의미한다. 즉 어떤 노드에 도착했어도 이를 탐색하지 않고 지나갈 수 있다. 그림으로 탐색 순서를 보이면 다음과 같다.

Inorder에서는 현재 노드가 노드 1일 때 왼쪽 자식 노드인 노드 2를 먼저 탐색하고 현재 노드인 노드 1, 마지막으로 오른쪽 자식 노드인 노드 3을 탐색한다.

Preorder에서는 현재 노드가 노드 1일 때 현재 노드인 노드 1을 먼저 탐색하고 왼쪽 자식 노드인 노드 2, 마지막으로 오른쪽 자식 노드인 노드 3을 탐색한다.

Postorder는 현재 노드가 노드 1일 때 왼쪽 자식 노드인 노드 2를 먼저 탐색하고 오른쪽 자식 노드인 노드 3, 마지막으로 현재 노드인 노드 1을 탐색한다.

Levelorder는 현재 노드가 노드 1일 때 현재 노드와 같은 높이의 노드를 모두 탐색하고 자식 노드로 넘어간다. 그리고 자식 노드와 같은 높이에 있는 모든 노드(노드 2, 노드 3)를 탐색하고 다시 자식 노드로 넘어간다. 결과적으로 같은 높이에 있는 노드 2, 노드 3의 자식 노드 노드 4, 노드 5, 노드 6, 노드 7은 순서대로 탐색되며 같은 높이, 레벨에 있는 노드들을 모두 탐색하기 때문에 Levelorder인 것이다.

 

이런 여러 가지 order 방식으로 탐색하려면 어떤 자료구조가 필요할까? 여기서는 Stack을 사용할 수 있다. 스택에서 노드를 하나씩 빼서 탐색한다고 가정할 때 맨 처음에는 최상위 노드를 스택에 삽입한 후 반복문으로 스택이 완전히 빌 때까지 하나씩 빼서 탐색하는 과정을 반복하는 것이다. Inorder 같은 경우 왼쪽, 현재, 오른쪽 순서대로 탐색해야 하기 때문에 각 노드를 탐색할 때마다 선입후출의 스택 특성상 오른쪽 자식 노드, 현재 노드, 왼쪽 자식 노드 순서대로 스택에 삽입하면 된다.

 

그런데 이때 트리가 더 깊다면 어떨까? 위의 그림에서는 보이지 않았지만 여러 레벨(2층 이상)의 트리 구조에서는 생각했던 것과 조금 다른 방식의 탐색을 볼 수 있다.

노드 1, 노드 2, 노드 3밖에 없을 때는 노드 2, 노드 1, 노드 3 순서대로 탐색했지만 노드 2가 자식 노드 4, 노드 5를 갖고 있을 때는 노드 1과 그의 자식 노드들을 탐색하는 게 아니라 노드 2의 자식 노드를 먼저 탐색하는 것이다. 즉 각각의 order에서 왼쪽 자식 노드나 오른쪽 자식 노드를 먼저 탐색한다는 것은 해당 자식 노드에도 자식 노드가 존재한다면 현재 노드를 탐색하는 대신 자식 노드의 자식 노드까지 더 깊게 들어가서 말단 노드까지 이동한다는 것이다.

 

그렇기 때문에 위의 그래프에서는 노드 4, 노드 2, 노드 5, 노드 1, 노드 3 순서대로 탐색하게 된다. 시작 노드 1의 왼쪽 자식 노드에는 노드 4, 노드 5가 자식으로 딸려있기 때문에 이를 더 탐색하고 자식이 더 이상 없는 말단 노드(leaf node)까지 도달했으면 그제야 탐색을 시작한다.

 

그런데 이렇게 제일 깊은 노드까지 들어갔다가 다시 나올 때는 무슨 일이 발생할 수 있을까? 스택을 활용하여 탐색할 노드를 왼쪽, 현재, 오른쪽 순서대로 스택에 삽입하면 노드 4까지 도달했을 때 스택에는 4, 2, 5, 1, 3 노드가 담겨있어야 스택에서 노드를 하나씩 빼면서 탐색하여 Inorder 방식으로 탐색할 수 있다. 그런데 말단 노드인 노드 4를 탐색한 후 자신의 부모 노드인 노드 2로 돌아가서 탐색할 때 자식 노드 중 이미 방문했던 노드인 노드 4를 기억하지 못한다면 다시 이를 스택에 삽입해서 무한히 해당 자식 노드를 탐색하는 문제가 발생할 수 있다. 이는 시간 초과 에러로 이어지기 때문에 방문한 자식 노드를 기억하기 위해 Set 자료구조를 활용할 수 있다.

import java.util.Stack;
import java.util.List;
import java.util.ArrayList;
import java.util.Set;
import java.util.HashSet;


class Solution {
    List<Integer> answer = new ArrayList<>();
    Stack<TreeNode> stack = new Stack<>();
    Set<TreeNode> visited = new HashSet<>();
    public List<Integer> inorderTraversal(TreeNode root) {
        if(root == null) return answer;
        stack.push(root);
        while(!stack.empty()){
            TreeNode parentNode = stack.pop();
            if(visited.contains(parentNode)){
                answer.add(parentNode.val);
                continue;
            }
             
            if(parentNode.right != null){
                stack.push(parentNode.right);
            }
            
            stack.push(parentNode);
            
            if(parentNode.left != null){
                stack.push(parentNode.left);
            }    
            
            visited.add(parentNode);
        }
        
        return answer;
    }
}

이를 이용하여 작성한 코드는 위와 같다. 먼저 스택에 최상위 노드(매개변수로 주어진 root)를 삽입하고 스택에서 노드를 하나씩 꺼내서 해당 노드의 왼쪽 자식 노드, 해당 노드, 해당 노드의 오른쪽 자식 노드를 스택에 삽입한다. 그리고 현재 노드는 방문한 노드로 표시하여 다음번에 다시 참조될 때 이미 방문한 노드라고 확인, 탐색 결과에 삽입하는 방식으로 구현하였다.

 

이미 방문한 노드일 경우 탐색 결과에 삽입하는 로직은 말단 노드에서도 자기 자신을 스택에 삽입하는 로직(stack.push(parentNode))을 기반으로 한 것이다. 위의 그래프에서 노드 2에서 노드 4, 노드 2, 노드 5를 스택에 삽입한 후 노드 4를 꺼내서 탐색할 때 노드 4는 말단 노드기 때문에 자기 자신만 다시 스택에 삽입한다. 그리고 다음 반복에서 스택에서 꺼낸 노드 4는 이미 방문한 노드로 visited에 기록되어 있기 때문에 탐색 결과에 삽입하고 그다음으로 꺼내는 노드 2 역시 방문한 노드로 기록되어 있으니 탐색 결과에 삽입한다. 오른쪽 자식 노드인 노드 5 역시 말단 노드기 때문에 노드 4처럼 자기 자신만 스택에 삽입 후 꺼내서 방문한 노드로 비교, 탐색 결과에 삽입한다. 이는 최상위 노드까지 반복되며 최상위 노드에서는 오른쪽 자식 노드부터 다시 비슷한 과정을 반복한다.

 

결과적으로 보면 최상위 노드에서 왼쪽 자식 노드(와 그의 서브 트리), 현재(최상위) 노드, 오른쪽 자식 노드(와 그의 서브 트리)를 순서대로 탐색하는 Inorder 방식으로 구현된 것이다. 각 자식 노드는 동일한 형태에 동일한 로직을 적용할 수 있기 때문에 재귀적으로 풀 수도 있지만 이번 튜토리얼에서는 Stack을 활용한 풀이를 원하고 있기 때문에 이처럼 작성하였다.

 

그렇지만 이게 메모리 상으로 효율적인 코드는 아니기 때문에 풀이 통계에서 스택을 사용한 반복(iterative) 방식의 풀이를 참고해보았는데 이는 다음과 같다.

class Solution {
    public List<Integer> inorderTraversal(TreeNode root) {
        List<Integer> result = new ArrayList<Integer>();
        if(root == null) {
            return result;
        }
        
        Stack<TreeNode> stack = new Stack<>();
        TreeNode current = root;
        while(!stack.isEmpty() || current != null) {
            while(current != null) {
                stack.push(current);
                current = current.left;
            }
            current =  stack.pop();
            result.add(current.val);
            current = current.right;
        }
        
        return result;
    }
}

훨씬 더 간단한 것을 볼 수 있는데 현재 탐색하고 있는 노드를 따로 변수로 선언하여 기억하고 있다. 탐색 시작 시 현재 노드에서 왼쪽 노드까지 최대한 이동하면서 스택에 노드를 삽입하는 것을 볼 수 있는데 이는 Inorder 방식으로 탐색하기 위한 방법이다. 그리고 끝까지 도달(말단 노드의 자식 노드, 즉 null)했다면 스택에서 노드를 하나 빼서 말단 노드를 참조한다. 이 말단 노드는 왼쪽 자식 노드로 쭉 내려온 노드기 때문에 Inorder에서 가장 먼저 탐색되어야 하는 노드기 때문에 이를 결과 리스트에 삽입하고 오른쪽 자식 노드를 탐색한다.

 

그런데 말단 노드는 자식 노드가 없기 때문에 오른쪽 자식 노드를 탐색하면 이는 null을 탐색하게 된다. 그래서 스택에서 노드를 하나 꺼내서 부모 노드로 이동한 후 부모 노드를 탐색(결과 리스트에 삽입)하고 오른쪽 자식 노드를 탐색하게 된다. 만약 오른쪽 자식 노드가 없다면 다시 스택에서 부모 노드를 꺼내서 점차 상위로 이동하면서 탐색하는 과정을 반복하게 된다.

 

결과적으로 왼쪽 자식 노드, 부모 노드, 오른쪽 자식 노드 순서대로 탐색하게 되기 때문에 Inorder 방식에 부합한다고 볼 수 있다. 부모 노드에서 오른쪽 자식 노드로 이동했을 때는 현재 노드가 null이 아니기 때문에 맨 처음에 수행했던 로직인 while 반복문(현재 노드를 스택에 삽입하고 자신의 왼쪽 자식 노드로 이동)을 반복하게 된다. 이때 노드가 null일 때까지 이동했다면, 또는 말단 노드라서 오른쪽 노드로 이동했더니 null이 됐다면 이동하기 전 스택에 삽입한 노드를 꺼내 탐색하기 때문에 결과적으로 오른쪽 자식 노드를 탐색한다.

 

즉 최대한(null이 될 때까지) 왼쪽 자식 노드로 파고들고 그 이후 스택에서 노드를 꺼내서 부모 노드, 그리고 오른쪽 자식 노드로 이동해서 다시 왼쪽 자식 노드로 파고드는 과정을 반복하면서 일관적인(왼쪽 자식 노드 우선) 방법으로 모든 노드들을 탐색하고 있는 것이다. 꽤나 효율적인 방법이라 인상 깊었다.

 

 

Inorder, preorder, postorder 같은 트리 탐색 방식은 예전에 학부 2학년 때 자료구조에서 배웠던 적이 있는데 지금 와서 다시 스택을 이용해 구현해보려고 하니 좀 까다로웠지만 복습이 많이 됐다. 트리 자료구조에도 좀 익숙해져야 할 것이다.