[Programmers] Trie 알고리즘

[3차] 자동완성
https://school.programmers.co.kr/learn/courses/30/lessons/17685

import java.io.*;
import java.util.*;

class Solution{
    // 노드 엔티티 
    class TrieNode{
        Map<Character, TrieNode> children;
        int count;
        Boolean isEnd;
        
        TrieNode() {
            this.children = new HashMap<>();
            this.count = 0;
            this.isEnd = false;
        }
    }
    // 트라이 자료구조
    class Trie{
        TrieNode root;
        
        Trie(){
            this.root = new TrieNode();
        }
        
        // 문자열을 통해 트라이 트리 만들기
        void addNode(String word){
            TrieNode node = root;
            for(char c : word.toCharArray()){
                
                // 해당 문자 노드 없으면 추가
                if(!node.children.containsKey(c)){
                    node.children.put(c,new TrieNode());
                }
                // 이동 후 count 추가
                node = node.children.get(c);
                node.count++;
            }
            // 마지막 노드에 도착했으면
            node.isEnd = true;
        }
        
        // 해당 문자열이 몇개를 검색해야 하는가? 
        int countNode(String word){
            TrieNode node = root;
            int count = 0;
            
            for(char c : word.toCharArray()){
                node = node.children.get(c);
                count++;
                // 이동하면서 1로 변하는 순간 or 끝까지 간다면 count 횟수 반환
                if(node.count == 1 || count == word.length()){
                    return count;
                }
            }
            return word.length();
        }
        
    }
    
    
    public int solution(String[] words) {
        int answer = 0;
        Trie trie = new Trie();
        // 추가
        for(String word : words){
            trie.addNode(word);
        }
        // 검색
        for(String word: words){
            answer += trie.countNode(word);
        }
        
        return answer;
    }
}

풀이

문제를 보고 트라이 자료구조를 떠올리게 되었다. 추가와 탐색 과정을 거치면서 최적화를 진행하고자 했다.

이건 트라이 자료구조를 통해서 go, gone, guild 단어를 넣게 되었을 때 어떻게 추가되는 지 그려보았다. 이렇게 트리가 완성되었다면 1로 바뀌거나 (자신만 남음), 단어 길이가 끝났을 때 몇번 검색(count)이 일어났는 지 확인해준다면 정답을 찾을 수 있다.