코딩 테스트와 알고리즘 문제 해결에서 효율적인 문자열 처리는 매우 중요합니다. 문자열 관련 문제는 많은 기업에서 자주 출제되기 때문에 이에 대한 이해는 필수적입니다. 오늘은 ‘트라이(Trie)’라는 자료구조를 활용하여 문자열을 효율적으로 저장하고 검색하는 방법에 대해 학습해 보겠습니다.

트라이(Trie)란?

트라이는 ‘프리픽스 트리(Prefix Tree)’ 또는 ‘디지털 트리(Digital Tree)’라고도 불리며, 문자열의 집합을 저장하기 위한 트리 기반 자료구조입니다. 주로 문자열 검색 또는 자동 완성 기능 구현에 적합합니다. 트라이의 주요 특징은 각 노드가 문자열의 각 문자에 대한 정보를 담고 있으므로, 공통 접두사를 공유하는 문자열에 대해 매우 효율적인 저장 및 검색이 가능하다는 점입니다. 예를 들어, 문자열 “hello”, “helicopter”, “helic”이 있을 경우, 첫 두 글자인 “he”를 공유하므로 이를 한 개의 노드로 표현할 수 있습니다.

문제 설명

다음은 트라이 자료구조를 사용할 수 있는 문제입니다:

문제: 주어진 단어 목록에서 특정 접두사를 가진 단어가 몇 개인지 구하는 함수를 구현하세요.

입력: 문자열 배열 words (ex: [“apple”, “app”, “apricot”, “banana”, “band”, “bandana”]) 및 문자열 prefix (예: “ap”)

출력: 주어진 prefix로 시작하는 단어의 개수

문제 해결 과정

1. 트라이 구현

먼저, 트라이 자료구조를 구현하기 위해 노드 클래스를 정의합니다. 각 노드는 다음과 같은 속성을 가집니다:

  • children: 다음 문자로 이동하기 위한 자식 노드들
  • isEndOfWord: 해당 노드가 단어의 끝인지 여부를 나타내는 불리언 변수
class TrieNode {
        Map children;
        boolean isEndOfWord;

        public TrieNode() {
            children = new HashMap<>();
            isEndOfWord = false;
        }
    }

2. 트라이 삽입 메서드

단어를 트리에 삽입하는 메서드를 구현합니다. 문자열의 각 문자를 순회하며 해당 문자의 자식 노드가 없다면 새로 생성하고, 마지막 문자에 도달했을 때는 isEndOfWordtrue로 설정합니다.

class Trie {
        private final TrieNode root;

        public Trie() {
            root = new TrieNode();
        }

        public void insert(String word) {
            TrieNode node = root;
            for (char ch : word.toCharArray()) {
                node = node.children.computeIfAbsent(ch, c -> new TrieNode());
            }
            node.isEndOfWord = true;
        }
    }

3. 접두사 검색 메서드

주어진 접두사로 시작하는 단어의 개수를 세기 위해, 먼저 접두사 노드를 찾아야 합니다. 접두사 노드를 찾은 후, 해당 노드에서 DFS를 통해 하위 트리의 단어 개수를 수셉니다.

public int countWordsWithPrefix(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toCharArray()) {
            if (!node.children.containsKey(ch)) {
                return 0; // 접두사가 존재하지 않으면 0 반환
            }
            node = node.children.get(ch);
        }
        return countWords(node);
    }

    private int countWords(TrieNode node) {
        int count = node.isEndOfWord ? 1 : 0; // 현재 노드가 단어의 끝이라면 1 추가
        for (TrieNode child : node.children.values()) {
            count += countWords(child); // 하위 노드에 대해 재귀 호출
        }
        return count;
    }

전체 코드

이제 위의 모든 메서드를 통합하여 완성된 코드 전체를 보여드리겠습니다.

import java.util.HashMap;
    import java.util.Map;

    class TrieNode {
        Map children;
        boolean isEndOfWord;

        public TrieNode() {
            children = new HashMap<>();
            isEndOfWord = false;
        }
    }

    class Trie {
        private final TrieNode root;

        public Trie() {
            root = new TrieNode();
        }

        public void insert(String word) {
            TrieNode node = root;
            for (char ch : word.toCharArray()) {
                node = node.children.computeIfAbsent(ch, c -> new TrieNode());
            }
            node.isEndOfWord = true;
        }

        public int countWordsWithPrefix(String prefix) {
            TrieNode node = root;
            for (char ch : prefix.toCharArray()) {
                if (!node.children.containsKey(ch)) {
                    return 0;
                }
                node = node.children.get(ch);
            }
            return countWords(node);
        }

        private int countWords(TrieNode node) {
            int count = node.isEndOfWord ? 1 : 0;
            for (TrieNode child : node.children.values()) {
                count += countWords(child);
            }
            return count;
        }
    }

    public class Main {
        public static void main(String[] args) {
            Trie trie = new Trie();
            trie.insert("apple");
            trie.insert("app");
            trie.insert("apricot");
            trie.insert("banana");
            trie.insert("band");
            trie.insert("bandana");
            
            System.out.println(trie.countWordsWithPrefix("ap")); // 출력: 3
        }
    }

결론

오늘은 트라이(Trie) 자료구조에 대해 살펴보았습니다. 트라이는 문자열 검색과 자동 완성 기능 구현에 매우 유용한 자료구조로, 효율적인 저장 및 검색이 가능합니다. 다양한 문자열 관련 문제를 해결하기 위해 트라이를 적극적으로 활용하면 좋겠습니다. 특히, 코딩 테스트를 준비하는 분들에게 트라이는 꼭 알아두어야 할 필수 자료구조입니다.

직접적인 문제 해결 뿐만 아니라, 트라이의 구조와 그 활용은 실무에서도 많이 사용되는 기법입니다. 트라이를 통해 시작하는 다양한 알고리즘 문제 풀이에 도전해 보시기 바랍니다!