백준 5052 전화번호 목록 python, java

Trie 에 대해 공부하고 관련 문제를 풀어 보았다.
이 문제의 경우 별도로 검색 과정은 필요없고 문자열 등록 과정에서 요구사항을 판단할 수 있었다.

📜 문제 바로 가기 : 전화번호 목록

제출코드

파이썬

import sys
input = sys.stdin.readline

from collections import defaultdict

class Node:
    def __init__(self, key, isWord=False):
        self.key = key
        self.isWord = isWord
        self.children = defaultdict(int)


class Trie:
    def __init__(self):
        self.head = Node(None)
    
    def insert(self, string):
        curr_node = self.head
        for char in string:
            if curr_node.children[char] == 0:
                curr_node.children[char] = Node(char)
            
            curr_node = curr_node.children[char]

            if curr_node.isWord == True:
                return False

        curr_node.isWord = True

        return True
  
t = int(input())

for i in range(t):
    n = int(input())
    trie = Trie()
    arr = []
    for _ in range(n):
        arr.append(input().rstrip())

    arr.sort(key=len)
    for string in arr:
        if trie.insert(string) == False:
            print("NO")
            break
    else:
        print("YES")

자바

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

public class Main {

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

        int T = Integer.parseInt(br.readLine());

        for (int tc = 0; tc < T; tc++) {
            int N = Integer.parseInt(br.readLine());
            String[] numberList = new String[N];
            for (int i = 0; i < N; i++) {
                String number = br.readLine();
                numberList[i] = number;
            }

            Arrays.sort(numberList, (s1, s2) -> s1.length() - s2.length());

            Trie trie = new Trie();
            boolean flag = true;
            for (String string : numberList) {
                if (trie.insert(string) == false) {
                    System.out.println("NO");
                    flag = false;
                    break;
                }
            }

            if (flag == true) {
                System.out.println("YES");
            }

        }
    }
}


class Node {
    Character key;
    boolean isWord;
    HashMap children;

    public Node(Character key) {
        this.key = key;
        this.isWord = false;
        this.children = new HashMap<Character, Node>();
    }
}

class Trie {
    Node head;

    public Trie() {
        this.head = new Node(null);
    }

    public boolean insert(String string) {
        Node curr_head = this.head;
        for (int i = 0; i < string.length(); i++) {
            Character cha = string.charAt(i);
            if (curr_head.children.containsKey(cha) == false) {
                Node newNode = new Node(cha);
                curr_head.children.put(cha, newNode);
            }

            curr_head = (Node) curr_head.children.get(cha); // (Node)

            if (curr_head.isWord == true) {
                return false;
            }
        }

        curr_head.isWord = true;
        return true;
    }

}

접근 방법

모든 전화번호를 길이순서로 오름차순 정렬해준다.
그리고 등록 과정을 거치는데 번호의 한 글자를 등록할 때마다 isWord 값이 True 인지 확인해준다.
번호 A를 등록하는 중에 어느 지점에서 isWord 값이 True 라면 번호 A에 연결되지 않고 다른 번호로 넘어 간다는 것을 의미하므로 False 를 리턴하여 처리해줬다.