Submission #1222319

#TimeUsernameProblemLanguageResultExecution timeMemory
1222319giorgi123glmSelling RNA Strands (JOI16_selling_rna)C++20
100 / 100
364 ms223228 KiB
#include <iostream>
#include <unordered_map>
#include <vector>
#include <algorithm>
#include <functional>
#include <array>
using namespace std;

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

vector <char> dehash = {
    'A', 'U', 'G', 'C'  
};

const int lim = 2e6 + 10;

struct trie {
    int free_node = 1;
    vector <array <int, 4> > tree;
    vector <int> str_node;
    vector <char> node_col;
    vector <int> in;
    vector <int> out;

    trie (int num_str) {
        tree.resize (lim + 1);
        str_node.resize (num_str + 1);
        node_col.resize (lim + 1);
        node_col[0] = '.'; 
    };

    void add (string str, int ind) {
        int node = 0;
        for (int i = 0; i < str.size(); i++) {
            if (tree[node][enhash[str[i]]] == 0) {
                node_col[free_node] = str[i];
                tree[node][enhash[str[i]]] = free_node++;
            }
            node = tree[node][enhash[str[i]]];
        }

        str_node[ind] = node;
    }

    int traverse (string str) {
        int node = 0;
        for (int i = 0; i < str.size(); i++) {
            if (tree[node][enhash[str[i]]] == 0)
                return -1;
            node = tree[node][enhash[str[i]]];
        }
        return node;
    }

    void calc_dfs () {
        in.clear();
        out.clear();
        in.resize (lim + 1);
        out.resize (lim + 1);
        int timer = 0;

        function <void(int,int)> dfs =
        [&](int p, int par)->void {
            in[p] = out[p] = ++timer;
            for (int& e : tree[p])
                if (e != par && e != 0)
                    dfs (e, p);
            out[p] = timer;
        };
        dfs (0, 0);
    }

    void debug () {
        function <void(int,int,int)> dfs =
        [&](int p, int par, int dep)->void {
            cout << string (dep, '\t') << node_col[p] << ' ' << in[p] << ' ' << out[p] << '\n';
            for (int& e : tree[p])
                if (e != par && e != 0)
                    dfs (e, p, dep + 1);
        };
        dfs (0, 0, 0);
    }
};

template <typename T>
class segment_tree {
    public:
        int siz = 1;
        vector <T> segtree;
        function <T(T, T)> merge;
        
        segment_tree(function <T(T, T)> IN) : merge (IN) {}

        void resize (int n) {
            while (siz < n)
                siz *= 2;
            segtree.resize (2 * siz);
        }

        void update (int ind, T K) {
            int u = siz + ind - 1;
            segtree[u] = merge (segtree[u], K);
            u /= 2;
            while (u)
                segtree[u] = merge (
                    segtree[2 * u], segtree[2 * u + 1]
                ), u /= 2;
        }

        T get (int L, int R, int u, int l, int r) {
            if (r < L || R < l)
                return T();

            if (L <= l && r <= R)
                return segtree[u];

            const int m = (l + r) / 2;
            return merge (
                get (L, R, 2 * u, l, m),
                get (L, R, 2 * u + 1, m + 1, r)
            );
        }

        T get (int L, int R) {
            return get (L, R, 1, 1, siz);
        }
};

int main () {
    int N = 0, Q = 0;
    cin >> N >> Q;

    vector <string> in (N + 1);
    for (int i = 1; i <= N; i++)
        cin >> in[i];

    trie preftrie (N + 1);
    trie sufftrie (N + 1);
    for (int i = 1; i <= N; i++) {
        preftrie.add (in[i], i);
        sufftrie.add ([](string str)->string {
            reverse (str.begin(), str.end());
            return str;
        }(in[i]), i);
    }
    preftrie.calc_dfs();
    sufftrie.calc_dfs();
    // preftrie.debug();
    // sufftrie.debug();

    vector <pair <int, int> > status (N + 1);
    for (int i = 1; i <= N; i++)
        status[i] = {
            preftrie.in[preftrie.str_node[i]],
            sufftrie.in[sufftrie.str_node[i]]
        };
    sort (status.begin() + 1, status.end());

    vector <array <int, 4> > q_in (Q + 1);
    for (int i = 1; i <= Q; i++) {
        string a = "", b = "";
        cin >> a >> b;

        reverse (b.begin(), b.end());

        int prefnode = preftrie.traverse (a);
        int suffnode = sufftrie.traverse (b);

        if (prefnode == -1 || suffnode == -1) {
            q_in[i] = {-1, -1, -1, -1};
            continue;
        }

        int& L1 = preftrie.in[prefnode], R1 = preftrie.out[prefnode];
        int& L2 = sufftrie.in[suffnode], R2 = sufftrie.out[suffnode];
        
        q_in[i] = {
            static_cast <int> (lower_bound (status.begin() + 1, status.end(), (pair <int, int>){L1, 0}) - status.begin()),
            static_cast <int> (lower_bound (status.begin() + 1, status.end(), (pair <int, int>){R1 + 1, 0}) - status.begin() - 1),
            L2, R2
        };
    }

    vector <vector <int> > add_to (lim + 1);
    for (int i = 1; i <= N; i++)
        add_to[status[i].second].emplace_back (i);

    vector <vector <array <int, 3> > > tasks (lim + 1);
    vector <int> res (2 * Q + 1);

    for (int i = 1, ind = 0; i <= Q; i++) {
        if (q_in[i][0] == -1)
            continue;

        tasks[q_in[i][2] - 1].push_back ({
            q_in[i][0], q_in[i][1], ind++
        });
        tasks[q_in[i][3]].push_back ({
            q_in[i][0], q_in[i][1], ind++
        });
    }

    segment_tree <int> segtree (
        [](int a, int b) {
            return a + b;
        }
    );
    segtree.resize (lim + 1);

    for (int i = 1; i <= lim; i++) {
        for (int& e : add_to[i])
            segtree.update (e, 1);
        for (array <int, 3>& e : tasks[i])
            res[e[2]] = segtree.get (e[0], e[1]);
    }

    for (int i = 1, ind = 0; i <= Q; i++) {
        if (q_in[i][0] == -1) {
            cout << 0 << '\n';
            continue;
        }

        cout << -res[ind] + res[ind + 1] << '\n';
        ind += 2;
    }
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...