[프로그래머스] 가사 검색 (C++)

 

프로그래머스

 

 

트라이

주어진 words로 트라이를 만들어준다.

struct Node {
    unordered_map<int, int> count;
    Node* next[26] = {};
};

 

 

count 는 해당 노드를 지나가는 word 들의 최대 사이즈의  (사이즈: 개수) 값을 가진다.

 

예를 들어, 트라이에 front를 추가한다면

f 노드의 count[5] += 1;

f -> r 노드의 count[5] += 1;

f -> r -> o 노드의 count[5] += 1;

이런 식이다.

 

next 는 다음 노드 정보이다.

 

 

frodo와 front가 있다면 다음과 같은 trie를 얻을 수 있다.

 

쿼리

쿼리는 세 가지 종류가 있다.

전체가 물음표로 된 경우와, 물음표 문자열을 접미사로 가지는 경우, 접두사로 가지는 경우이다.

 

 

전체가 물음표인 경우

전체가 물음표인 쿼리의 경우, 트라이 루트노드의 count 에서 쿼리 사이즈를 찾으면 답을 구할 수 있다.

 

 

물음표 문자열을 접미사로 가지는 경우

예를 들어, 위 트라이에서 fro?? 를 찾게 되면 f -> r -> o 까지 이동한 다음, count 에서 인덱스 5('fro??'의 사이즈)에 접근하면 2가 나오게 된다.

 

따라서 물음표 이전까지 이동한 다음 count 에서 쿼리 사이즈의 인덱스에 접근하면 답을 구할 수 있다.

 

 

물음표 문자열을 접두사로 가지는 경우

물음표 문자열을 접두사로 가지는 경우엔, 단어를 뒤집은 상태로 만든 트라이가 필요하다.

 

frodo와 front가 있다면 다음과 같이 나타낼 수 있다.

 

 

 

예를 들어 ??ont 를 찾는 경우에, 뒤에서부터 t -> n -> o 까지 이동한 다음, 마찬가지로 count 에서 5의 인덱스에 접근하면 1이 나옴을 알 수 있다.

 

 

예시 입력 트라이 그림

 

 

 

 

정답 코드

#include <bits/stdc++.h>
using namespace std;

struct Node {
    unordered_map<int, int> count;
    Node* next[26] = {};
};

vector<int> solution(vector<string> words, vector<string> queries) {
    Node* g = new Node();
    Node* r_g = new Node();

    for (const string& word : words) {
        Node* node = g;
        Node* r_node = r_g;

        int sz = word.size();
        g->count[sz]++;
        r_g->count[sz]++;

        for (int i = 0; i < sz; i++) {
            int value = word[i] - 'a';
            int r_value = word[sz - i - 1] - 'a';

            if (!node->next[value]) {
                node->next[value] = new Node();
            }

            if (!r_node->next[r_value]) {
                r_node->next[r_value] = new Node();
            }

            node = node->next[value];
            node->count[sz]++;

            r_node = r_node->next[r_value];
            r_node->count[sz]++;
        }
    }

    vector<int> ans;
    for (const string& q : queries) {
        int sz = q.size();
        int result = 0;

        if (q.front() == '?' && q.back() == '?') {
            result = g->count[sz];
        } else if (q.front() == '?') {
            Node* node = r_g;
            for (int i = sz - 1; i >= 0 && node; i--) {
                if (q[i] == '?') {
                    result = node->count[sz];
                    break;
                }
                node = node->next[q[i] - 'a'];
            }
        } else if (q.back() == '?') {
            Node* node = g;
            for (int i = 0; i < sz && node; i++) {
                if (q[i] == '?') {
                    result = node->count[sz];
                    break;
                }
                node = node->next[q[i] - 'a'];
            }
        }
        ans.push_back(result);
    }
    return ans;
}