백준16780 Selling RNA Strands

문제 링크

  • http://icpc.me/16780

문제 출처

  • 2016 JOI Open Contest 2번

사용 알고리즘

  • Trie

풀이

문자열 $P, Q$가 쿼리로 주어지는데, $P$를 Prefix로 갖는 문자열들의 목록은 주어진 문자열 $S_i$를 정렬한 뒤 lower_bound/upper_bound를 이용해 구할 수 있습니다. 이렇게 해서 구한 문자열 중 $Q$를 Suffix를 갖는 문자열을 구하면 됩니다.

$S_i$들의 Suffix Trie를 구축합시다. $Q$를 Suffix로 갖는 문자열의 목록은 Trie에서 $Q$를 찾은 뒤, 마지막에 방문한 정점을 포함하는 문자열입니다. 그러므로 Trie의 각 정점에 해당 정점을 포함하는 “문자열의 인덱스”를 정렬해서 들고 있으면, lower_bound/upper_bound를 이용해 $P$를 Prefix, $Q$를 Suffix로 갖는 문자열의 개수를 구할 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
/******************************
Author: jhnah917(Justice_Hui)
g++ -std=c++17 -DLOCAL
******************************/

#include <bits/stdc++.h>
#define x first
#define y second
#define all(v) v.begin(), v.end()
#define compress(v) sort(all(v)), v.erase(unique(all(v)), v.end())
#define IDX(v, x) lower_bound(all(v), x) - v.begin()
using namespace std;

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

unordered_map<char, int> ID = {
    {'A',0},
    {'C',1},
    {'G',2},
    {'U',3}
};

struct Trie{
    Trie *ch[4];
    vector<int> val;
    Trie(){ memset(ch, 0, sizeof ch); }
    void Insert(const char *s, int id){
        val.push_back(id); if(!*s) return;
        if(!ch[ID[*s]]) ch[ID[*s]] = new Trie;
        ch[ID[*s]]->Insert(s+1, id);
    }
    int Query(const char *s, int l, int r){
        if(!*s) return lower_bound(all(val), r) - lower_bound(all(val), l);
        if(ch[ID[*s]]) return ch[ID[*s]]->Query(s+1, l, r);
        return 0;
    }
} *Root;

int N, Q;
string A[101010], B[101010];

int main(){
    ios_base::sync_with_stdio(false); cin.tie(nullptr);
    cin >> N >> Q; Root = new Trie;
    for(int i=1; i<=N; i++) cin >> A[i];
    sort(A+1, A+N+1);
    copy(A+1, A+N+1, B+1);
    for(int i=1; i<=N; i++) reverse(all(B[i]));
    for(int i=1; i<=N; i++) Root->Insert(B[i].c_str(), i);

    for(int i=1; i<=Q; i++){
        string P, S; cin >> P >> S; reverse(all(S));
        int L = lower_bound(A+1, A+N+1, P) - A;
        P[P.size()-1]++;
        int R = upper_bound(A+1, A+N+1, P) - A;
        cout << Root->Query(S.c_str(), L, R) << "\n";
    }
}