백준 19585번

오늘 풀어본 문제는 백준의 19585번 문제¹이다. 문제 풀이에 사용한 언어는 C++ 이다.

solved.ac 기준 CLASS

CLASS 6

문제 정보

이 문제의 내용과 조건은 다음과 같다.

문제

Sogang ICPC Team에는 색상 이름과 닉네임의 순서로 이여서 팀명을 지으면 ICPC 리저널에서 수상할 수 있다는 전설이 있다. 색상 이름들과 닉네임들이 주어질 때, $Q$ 개의 팀에 대해 다음 리저널에서 수상할 수 있을지 전설에 기반해 알려주는 프로그램을 작성하자.

입력

첫째 줄에는 색상의 종류 $C$ 와 닉네임의 개수 $N$ 이 주어진다. $(1 \le C, N \le 4,000)$

다음 $C$ 개의 줄에는 색상 이름 $C$ 개가 한 줄에 하나씩 주어진다. 색상 이름은 중복되지 않는다.

다음 $N$ 개의 줄에는 Sogang ICPC Team 구성원들의 닉네임 $N$ 개가 한 줄에 하나씩 주어진다. 닉네임도 중복되지 않는다.

다음 줄에는 팀의 개수 $Q$ 가 주어진다. $(1 \le Q \le 20,000)$

다음 $Q$ 개의 줄에는 팀명 $Q$ 개가 한 줄에 하나씩 주어진다. 팀명은 중복되지 않는다.

모든 색상 이름의 길이와 닉네임의 길이는 $1,000$ 글자를 넘지 않는다. 모든 팀명은 $2,000$ 글자를 넘지 않는다. 모든 문자열은 알파벳 소문자로만 이루어져 있다.

출력

팀명이 입력된 순서대로, 전설에 따라 해당 팀이 다음 리저널에서 수상할 수 있다면 Yes, 아니라면 No를 출력한다.

풀이과정

1번째 시도

문제를 해결하기 위해 사용한 방식은 다음과 같다.

Trie 를 구현한다.
Trie 를 이용하여 각각 색상과 닉네임을 저장한다.
색상이 저장된 Trie 로부터 입력받은 팀의 이름에서 닉네임이 시작될 수 있는 위치의 후보를 얻어내고, 그 위치부터 추출된 문자열이 존재하는 닉네임인지 확인한다.

(닉네임이 시작될 수 있는 위치가 여러 개가 될 수 있는 이유는, 색상의 이름이 yellowishorange 처럼 다른 색상의 이름을 접두어로 가질 수 있기 때문이다.)
상황에 맞게 Yes 나 No 를 출력한다.

코드는 다음과 같이 작성하였다. 제출할 때는 헤더파일의 코드를 붙여넣어 제출하였다.

trie.hpp

#ifndef __TRIE_HPP__
#define __TRIE_HPP__

#include <string>
#include <vector>

struct TrieNode {
    bool endOfWord;
    std::vector<int> next;
    int nextLetters;

    TrieNode() : endOfWord(false), nextLetters(0) {
        next.resize(26, -1);
    }
};

struct Trie {
    std::vector<TrieNode> nodes;

    Trie() {
        nodes.emplace_back();
    }

    void insert(const std::string& str) {
        int curr = 0;
        for (char c : str) {
            int idx = c - 'a';
            if (nodes[curr].next[idx] == -1) {
                nodes[curr].nextLetters++;

                nodes[curr].next[idx] = nodes.size();
                nodes.emplace_back();
            }
            curr = nodes[curr].next[idx];
        }
        nodes[curr].endOfWord = true;
    }

    bool find(const std::string& str) {
        int curr = 0;
        for (char c : str) {
            int idx = c - 'a';

            if (nodes[curr].next[idx] == -1) {
                return false;
            }
            else {
                curr = nodes[curr].next[idx];
            }
        }
        return nodes[curr].endOfWord;
    }

    std::vector<int> prefixSearch(const std::string& str) {
        int curr = 0;
        std::vector<int> result;

        for (int i=0; i<str.size(); i++) {
            int idx = str[i] - 'a';

            if (nodes[curr].endOfWord) {
                result.emplace_back(i);
            }

            if (nodes[curr].next[idx] == -1) {
                return result;
            }
            else {
                curr = nodes[curr].next[idx];
            }
        }

        return result;
    }
};

#endif // !__TRIE_HPP__

main.cpp

#include <bits/stdc++.h>
#include "trie.hpp"

using namespace std;

using ll = long long;
using ull = unsigned long long;

using pii = pair<int, int>;
using pll = pair<ll, ll>;

int main(void) {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int C, N;
    cin >> C >> N;

    Trie colorTrie;

    for (int i=0; i<C; i++) {
        string color;
        cin >> color;
        colorTrie.insert(color);
    }

    Trie nicknameTrie;

    for (int i=0; i<N; i++) {
        string nickname;
        cin >> nickname;
        nicknameTrie.insert(nickname);
    }

    int Q;
    cin >> Q;

    for (int i=0; i<Q; i++) {
        string team;
        cin >> team;

        vector<int> colorPrefixEndingIndex = colorTrie.prefixSearch(team);
        bool found = false;

        for (auto& nickStart : colorPrefixEndingIndex) {
            if (nicknameTrie.find(team.substr(nickStart))) {
                found = true;
                break;
            }
        }

        if (found) {
            cout << "Yes\n";
        }
        else {
            cout << "No\n";
        }
    }

    return 0;
}

실행 결과 ‘메모리 초과’ 가 떴다.

2번째 시도

메모리 초과를 해결하기 위해 마지막에 닉네임 일치여부를 확인하는 것은 해시테이블을 이용하기로 했다.

코드는 다음과 같이 수정하였고, main.cpp 파일만 수정하였다.

main.cpp

#include <bits/stdc++.h>
#include <unordered_set>
#include "trie.hpp"

using namespace std;

using ll = long long;
using ull = unsigned long long;

using pii = pair<int, int>;
using pll = pair<ll, ll>;

int main(void) {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int C, N;
    cin >> C >> N;

    Trie colorTrie;

    for (int i=0; i<C; i++) {
        string color;
        cin >> color;
        colorTrie.insert(color);
    }

    unordered_set<string> nicknameTable;

    for (int i=0; i<N; i++) {
        string nickname;
        cin >> nickname;
        nicknameTable.insert(nickname);
    }

    int Q;
    cin >> Q;

    for (int i=0; i<Q; i++) {
        string team;
        cin >> team;

        vector<int> colorPrefixEndingIndex = colorTrie.prefixSearch(team);
        bool found = false;

        for (auto& nickStart : colorPrefixEndingIndex) {
            if (nicknameTable.find(team.substr(nickStart)) != nicknameTable.end()) {
                found = true;
                break;
            }
        }

        if (found) {
            cout << "Yes\n";
        }
        else {
            cout << "No\n";
        }
    }

    return 0;
}

그러자 모든 테스트 케이스를 통과하고 ‘맞았습니다’가 나오는 것을 확인할 수 있었다.

마무리

이번에 풀었던 문제는 꽤 흥미로운 요소가 많았다. 새로운 자료구조인 Trie 의 사용도 흥미롭지만, 문제를 푸는 과정에서 끄투에서나 쓰던 yellowishorange 를 떠올려 실수를 미리 제거할 수 있었던 건 정말 흥미로운 경험이었다. 역시 쓸모없는 경험이란건 없는걸까?

내가 만든 Trie 자료구조는 내 레포지토리²에서 확인할 수 있다. 필요한 사람은 참고하길 바란다.

오늘의 PS는 여기까지!

1: https://www.acmicpc.net/problem/19585
2: https://github.com/ChoiCube84/baekjoon-solutions/blob/main/C%2B%2B/custom_data_structures/trie.hpp

← Previous Post Next Post →