2020 - 가사 검색

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

코딩테스트 연습 - 가사 검색

[ 문제풀이 ]

Trie 자료구조를 활용하는 문제이다. 모든 문자열에 대해서 정방향, 역방향 Trie를 구성해 준 후 쿼리를 해주면 된다.

​각각 한 개의 Trie를 통해 문자열을 관리해줘서 해결할 수도 있지만 문자열의 길이에 따라 쿼리 값이 결정되기 때문에 조금 구현하기가 난해할 것이다.

​이를 위해 모든 문자열의 길이만큼 Trie 자료구조를 구현해준다. 문자열이 최대 10000이기 때문에 2만 개의 Trie를 구현해준다.

​이렇게 구현한 후 쿼리를 하게 되면 "????" 형태의 모든 문자열이 '?'인 경우만 별도로 처리해주면 해결할 수 있다. 모든 문자열이 '?'인 경우는 26개의 자식 수를 더해주면 된다.

class Solution {
	class Trie {
		char data;
		int childCnt;
		Trie[] child;

		public Trie() {
			child = new Trie[26];
		}

        //첫번째 문자부터 마지막 문자까지 진행
		public void createTrie(Trie root, int idx, String word){
			if (idx == word.length()) return;

			char now = word.charAt(idx);
			int pos = (int)now - 97;
			
			if (root.child[pos] == null) {
				root.child[pos] = new Trie();
				root.child[pos].data = now;
			}

			root.child[pos].childCnt++;
			createTrie(root.child[pos], idx + 1, word);
		}

        //마지막 문자부터 첫번째 문자까지 진행
		public void createReverseTrie(Trie root, int idx, String word) {
			if (idx < 0) return;

			char now = word.charAt(idx);
			int pos = (int)now - 97;
			
			if (root.child[pos] == null) {
				root.child[pos] = new Trie();
				root.child[pos].data = now;
			}

			root.child[pos].childCnt++;
			createReverseTrie(root.child[pos], idx - 1, word);
		}

		public int queryA(Trie root, int idx, String word) {
			if (root == null) return 0;

			char now = word.charAt(idx);
			int pos = (int)now - 97;

			if (now == '?') {
				return root.childCnt;
			}

			return queryA(root.child[pos], idx + 1, word);
		}

		public int queryB(Trie root, int idx, String word) {
			if (root == null) return 0;

			char now = word.charAt(idx);
			int pos = (int)now - 97;
			
			if (now == '?') {
                //모든 문자열이 '?'로 구성된 경우 ('?'가 첫번째인 경우)
				if (idx == word.length() - 1) {
					int sum = 0;
					for (int i = 0; i < 26; ++i) {
						if (root.child[i] == null) continue;
						sum += root.child[i].childCnt;
					}
					return sum;
				}
				return root.childCnt;
			}

			return queryB(root.child[pos], idx - 1, word);
		}
	}

	public int[] solution(String[] words, String[] queries) {
		int[] answer = new int[queries.length];
		
		Trie[][] root = new Trie[2][10001];
		for (int i = 0; i < 2; ++i) {
			for (int j = 0; j < 10001; ++j) {
				root[i][j] = new Trie();
			}
		}

        //모든 문자열으로 정방향, 역방향 trie 생성
		for (String str : words) {
			int len = str.length();
			root[0][len].createTrie(root[0][len], 0, str);
			root[1][len].createReverseTrie(root[1][len], len - 1, str);
		}

        //query
		for (int i = 0; i < queries.length; ++i) {
			int len = queries[i].length();
			if (queries[i].charAt(0) == '?') {
				answer[i] = root[1][len].queryB(root[1][len], len - 1, queries[i]);
			} else {
				answer[i] = root[0][len].queryA(root[0][len], 0, queries[i]);
			}
		}

		return answer;
	}
}