[프로그래머스] 가사 검색

문제 설명

https://school.programmers.co.kr/learn/courses/30/lessons/60060

문제 풀이

• 트라이를 이용하여 풀이하였다.
• 일반적인 트라이를 사용하되 단어의 정순, 역순 트라이를 각각 생성하여 가사를 insert할 때 각각 insert 해준다.
• 쿼리 개수는 ?로 시작하면 역순 트라이, ?로 끝나면 정순 트라이에서 아래와 같은 로직으로 계산한다.

  1. ?이전 문자까지 트라이 노드 순회를 통하여 진행한다.
  2. DFS를 통해서 해당 노드 하위의 남은 길이만큼의 길이를 가지는 모든 단어의 개수를 계산한다.

소스 코드 (JAVA)

import java.util.*;
class Solution {

    public class Node {
        Map<Character, Node> child = new HashMap<>();
        Boolean isEnd = false;
    }

    public class Trie {
        Node root = new Node();

        public int caculateCnt (Node cur, int remain) {
            //결과 리턴
            if (remain == 0) {
                if (cur.isEnd) return 1;
                else return 0;
            }

            int result = 0;
            for (Character key : cur.child.keySet()) {
                result += caculateCnt(cur.child.get(key), remain - 1);
            }
            return result;
        }

        public void insert (String input) {
            Node node = this.root;

            for (int i = 0; i < input.length(); i++)
                node = node.child.computeIfAbsent(input.charAt(i), k -> new Node());

            node.isEnd = true;
        }

        public int matchCnt (String input) {
            Node node = this.root;

            for (int i = 0; i < input.length(); i++) {
                //결과 계산 시작
                if (input.charAt(i) == '?') {
                    return caculateCnt(node, input.length() - i);
                }
                //? 이전까지 진행
                node = node.child.getOrDefault(input.charAt(i), null);

                if (node == null)
                    break;
            }
            return 0;
        }
    }

    public int[] solution(String[] words, String[] queries) {
        Trie forwardTrie = new Trie();
        Trie backwardTrie = new Trie();

        //트라이 insert 정순, 역순
        StringBuffer sb;
        for (int i = 0; i < words.length; i++) {
            sb = new StringBuffer(words[i]);
            forwardTrie.insert(words[i]);
            backwardTrie.insert(sb.reverse().toString());
        }

        //쿼리별 개수리턴
        Map<String, Integer> resultMap = new HashMap<>();
        int[] result = new int[queries.length];
        for (int i = 0; i < queries.length; i++) {
            if (resultMap.containsKey(queries[i])) {
                result[i] = resultMap.get(queries[i]);
            } else {
                sb = new StringBuffer(queries[i]);
                if (queries[i].charAt(0) != '?')
                    result[i] = forwardTrie.matchCnt(queries[i]);
                else
                    result[i] = backwardTrie.matchCnt(sb.reverse().toString());
                resultMap.put(queries[i], result[i]);
            }
        }

        return result;
    }
}