Submission #1279916

#	Time	Username	Problem	Language	Result	Execution time	Memory
1279916		hotranminhky0405_	Selling RNA Strands (JOI16_selling_rna)	C++17	100 / 100	713 ms	42336 KiB

selling_rna

// Compile with: g++ -std=gnu++17 -O2 -pipe -static -s -o selling_rna selling_rna.cpp
#include <bits/stdc++.h>
using namespace std;

int chmap(char c){
    if(c=='A') return 0;
    if(c=='G') return 1;
    if(c=='C') return 2;
    return 3; // 'U'
}

struct Trie {
    struct Node {
        int next[4];
        int end_id; // id for a pattern endpoint, -1 if none
        Node(){ fill(begin(next), end(next), -1); end_id = -1; }
    };
    vector<Node> nodes;
    Trie(){ nodes.emplace_back(); }
    int add_pattern(const string &s, int pat_id){
        int cur = 0;
        for(char c: s){
            int x = chmap(c);
            if(nodes[cur].next[x] == -1){
                nodes[cur].next[x] = nodes.size();
                nodes.emplace_back();
            }
            cur = nodes[cur].next[x];
        }
        // if multiple patterns identical, keep same end_id (we will map many queries to same pattern id)
        if(nodes[cur].end_id == -1) nodes[cur].end_id = pat_id;
        return nodes[cur].end_id;
    }
    // traverse string s, and for every node encountered that has end_id != -1, call visitor(end_id)
    template<typename F>
    void visit_prefixes_in_string(const string &s, F visitor){
        int cur = 0;
        for(char c: s){
            int x = chmap(c);
            if(nodes[cur].next[x] == -1) return;
            cur = nodes[cur].next[x];
            if(nodes[cur].end_id != -1){
                visitor(nodes[cur].end_id);
            }
        }
    }
    // for reversed trie, traverse reversed string
    template<typename F>
    void visit_suffixes_in_reversed_string(const string &rev_s, F visitor){
        int cur = 0;
        for(char c: rev_s){
            int x = chmap(c);
            if(nodes[cur].next[x] == -1) return;
            cur = nodes[cur].next[x];
            if(nodes[cur].end_id != -1){
                visitor(nodes[cur].end_id);
            }
        }
    }
};

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int N, M;
    if(!(cin >> N >> M)) return 0;
    vector<string> S(N);
    for(int i=0;i<N;i++) cin >> S[i];

    vector<pair<string,string>> queries(M);
    vector<string> Ps(M), Qs(M);
    for(int j=0;j<M;j++){
        cin >> Ps[j] >> Qs[j];
        queries[j] = {Ps[j], Qs[j]};
    }

    // Build unique pattern lists for P and Q (we want to map identical P's to same id)
    unordered_map<string,int> pid_map;
    unordered_map<string,int> qid_map;
    pid_map.reserve(M*2);
    qid_map.reserve(M*2);
    int pid_cnt = 0, qid_cnt = 0;
    vector<int> q_to_pid(M), q_to_qid(M);

    for(int j=0;j<M;j++){
        auto &p = Ps[j];
        auto &q = Qs[j];
        auto itp = pid_map.find(p);
        if(itp==pid_map.end()){
            pid_map[p] = pid_cnt++;
        }
        auto itq = qid_map.find(q);
        if(itq==qid_map.end()){
            qid_map[q] = qid_cnt++;
        }
    }

    // Build tries
    Trie preTrie, sufTrie;
    // map pattern id -> trie end id used in trie (we maintain consistent ids equal to 0..pid_cnt-1 and 0..qid_cnt-1)
    // We'll add patterns to trie in arbitrary order but ensure end_id stored equals the pattern id we assigned.
    // So create vector of patterns by id:
    vector<string> patternsP(pid_cnt);
    vector<string> patternsQ(qid_cnt);
    for(auto &kv : pid_map) patternsP[kv.second] = kv.first;
    for(auto &kv : qid_map) patternsQ[kv.second] = kv.first;

    // add to tries
    for(int i=0;i<pid_cnt;i++){
        preTrie.add_pattern(patternsP[i], i);
    }
    for(int i=0;i<qid_cnt;i++){
        string rq = patternsQ[i];
        reverse(rq.begin(), rq.end());
        sufTrie.add_pattern(rq, i);
    }

    // For each pattern id, maintain vector<int> of S indices that have that prefix/suffix
    vector<vector<int>> listP(pid_cnt), listQ(qid_cnt);

    // Fill lists: for each S[i], traverse preTrie along S and add i to endpoints
    for(int i=0;i<N;i++){
        preTrie.visit_prefixes_in_string(S[i], [&](int pid){
            listP[pid].push_back(i);
        });
        string revS = S[i];
        reverse(revS.begin(), revS.end());
        sufTrie.visit_suffixes_in_reversed_string(revS, [&](int qid){
            listQ[qid].push_back(i);
        });
    }

    // The vectors are already in increasing order of i (we traversed i from 0..N-1).
    // Prepare mapping from query j to pid/qid
    for(int j=0;j<M;j++){
        q_to_pid[j] = pid_map[Ps[j]];
        q_to_qid[j] = qid_map[Qs[j]];
    }

    // Cache for computed pair answers
    unordered_map<unsigned long long,int> cache;
    cache.reserve(M*2);

    // Answer queries
    // For each query, take smaller vector and binary_search into larger vector to count intersection
    for(int j=0;j<M;j++){
        int a = q_to_pid[j];
        int b = q_to_qid[j];
        unsigned long long key = ((unsigned long long)a << 32) | (unsigned long long)(unsigned)b;
        auto itc = cache.find(key);
        if(itc != cache.end()){
            cout << itc->second << '\n';
            continue;
        }
        auto &A = listP[a];
        auto &B = listQ[b];
        int ans = 0;
        // iterate smaller
        if(A.size() <= B.size()){
            for(int x: A){
                // binary search x in B
                if(binary_search(B.begin(), B.end(), x)) ++ans;
            }
        }else{
            for(int x: B){
                if(binary_search(A.begin(), A.end(), x)) ++ans;
            }
        }
        cache[key] = ans;
        cout << ans << '\n';
    }
    return 0;
}

• Subtask 1: 10 / 10
• Subtask 2: 25 / 25
• Subtask 3: 25 / 25
• Subtask 4: 40 / 40