Submission #742156

#	Time	Username	Problem	Language	Result	Execution time	Memory
742156		GrindMachine	Selling RNA Strands (JOI16_selling_rna)	C++17	100 / 100	848 ms	781128 KiB

selling_rna

// Om Namah Shivaya

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x, y) ((x + y - 1) / (y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl

#define rep(i, n) for(int i = 0; i < n; ++i)
#define rep1(i, n) for(int i = 1; i <= n; ++i)
#define rev(i, s, e) for(int i = s; i >= e; --i)
#define trav(i, a) for(auto &i : a)

template<typename T>
void amin(T &a, T b) {
    a = min(a, b);
}

template<typename T>
void amax(T &a, T b) {
    a = max(a, b);
}

#ifdef LOCAL
#include "debug.h"
#else
#define debug(x) 42
#endif

/*

refs:
https://github.com/Yehezkiel01/CompetitiveProgramming/blob/master/JOIOC/JOIOC-16-selling_rna.cpp


count #of strings that have prefix = p, suffix = q
if suffix condition doesn't exist, then we can solve this using a simple trie

how to consider this condition

we can add all input strings to a trie

for a query, we can go to the node representing the string p in the trie
prefix condition is satisfied

we want to count the #of leaves in the subtree of this node that correspond to q

put reversed string @ every leaf

for all (reversed) strings in the subtree of the current node, find the #of strings that are a prefix of reversed q

consider the original trie as a tree
store a trie @ every node and do small to large merging

*/

const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

int get(char ch) {
    int x = -1;
    if (ch == 'A') x = 0;
    if (ch == 'G') x = 1;
    if (ch == 'C') x = 2;
    if (ch == 'U') x = 3;

    return x;
}

vector<int> ans(N);

struct Trie {
    struct node {
        int f[4];
        int cnt, end;
        string str = "";

        node() {
            memset(f, -1, sizeof f);
            cnt = 0, end = 0;
            str = "";
        }
    };

    vector<node> tr;
    int siz = 0;

    Trie() {
        tr.pb(node());
        siz++;
    }

    void insert(string s, int add) {
        int u = 0;
        trav(ch, s) {
            int x = get(ch);

            if (tr[u].f[x] == -1) {
                tr.pb(node());
                tr[u].f[x] = siz++;
            }

            u = tr[u].f[x];
            tr[u].cnt += add;
        }

        tr[u].end++;

        reverse(all(s));
        tr[u].str = s;
    }

    vector<vector< pair<string, int> >> queries;

    void init_queries() {
        queries.resize(siz);
    }

    void insert_query(string p, string q, int id) {
        int u = 0;

        trav(ch, p) {
            int x = get(ch);
            if (tr[u].f[x] == -1) {
                u = -1;
                break;
            }

            u = tr[u].f[x];
        }

        if (u != -1) {
            reverse(all(q));
            queries[u].pb({q, id});
        }
    }

    void merge(Trie &trie1, Trie &trie2, int u1, int u2) {
        rep(x, 4) {
            int v1 = trie1.tr[u1].f[x];
            int v2 = trie2.tr[u2].f[x];

            if (v2 == -1) conts;

            if (v1 == -1) {
                trie1.tr.pb(node());
                trie1.tr[u1].f[x] = trie1.siz++;
                v1 = trie1.tr[u1].f[x];
            }

            trie1.tr[v1].cnt += trie2.tr[v2].cnt;
            merge(trie1, trie2, v1, v2);
        }
    }

    vector<Trie> trie;

    void dfs(int u) {
        rep(x, 4) {
            int v = tr[u].f[x];
            if (v == -1) conts;

            dfs(v);

            if (trie[v].siz > trie[u].siz) {
                swap(trie[u], trie[v]);
            }

            merge(trie[u], trie[v], 0, 0);
        }

        if (!tr[u].str.empty()) {
            trie[u].insert(tr[u].str, tr[u].end);
        }

        for (auto [q, id] : queries[u]) {
            int curr = 0;

            trav(ch, q) {
                int x = get(ch);

                if (trie[u].tr[curr].f[x] == -1) {
                    curr = -1;
                    break;
                }

                curr = trie[u].tr[curr].f[x];
            }

            if (curr != -1) {
                ans[id] = trie[u].tr[curr].cnt;
            }
        }
    }

    void dfs() {
        trie.resize(siz);
        dfs(0);
    }
};

void solve(int test_case)
{
    int n, m; cin >> n >> m;
    vector<string> a(n);
    rep(i, n) cin >> a[i];

    // insert all initial strings into trie
    Trie trie;
    rep(i, n) trie.insert(a[i], 1);

    // insert all query strings into their appropriate positions
    trie.init_queries();

    rep(i, m) {
        string p, q; cin >> p >> q;
        trie.insert_query(p, q, i);
    }

    // answer queries by dfs + small to large merging
    trie.dfs();
    rep(i, m) cout << ans[i] << endl;
}

int main()
{
    fastio;

    int t = 1;
    // cin >> t;

    rep1(i, t) {
        solve(i);
    }

    return 0;
}

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