Codeforces 1252D Find String in a Grid SA + BIT

#include<bits/stdc++.h>
using namespace std;

const int N = (int)7e5 + 7;
const int LOG = 20;

int Log[N];
struct SA {
int sa[N], rk[N], ht[N], s[N << 1], t[N << 1], p[N], cnt[N], cur[N];
int rmq[N][LOG];
#define pushS(x) sa[cur[s[x]]--] = x
#define pushL(x) sa[cur[s[x]]++] = x
#define inducedSort(v) \
    fill_n(sa, n, -1); fill_n(cnt, m, 0);                                     \
    for (int i = 0; i < n; i++) cnt[s[i]]++;                                  \
    for (int i = 1; i < m; i++) cnt[i] += cnt[i-1];                           \
    for (int i = 0; i < m; i++) cur[i] = cnt[i]-1;                            \
    for (int i = n1-1; ~i; i--) pushS(v[i]);                                  \
    for (int i = 1; i < m; i++) cur[i] = cnt[i-1];                            \
    for (int i = 0; i < n; i++) if (sa[i] > 0 &&  t[sa[i]-1]) pushL(sa[i]-1); \
    for (int i = 0; i < m; i++) cur[i] = cnt[i]-1;                            \
    for (int i = n-1;  ~i; i--) if (sa[i] > 0 && !t[sa[i]-1]) pushS(sa[i]-1);
void sais(int n, int m, int *s, int *t, int *p) {
    int n1 = t[n-1] = 0, ch = rk[0] = -1, *s1 = s+n;
    for (int i = n-2; ~i; i--) t[i] = s[i] == s[i+1] ? t[i+1] : s[i] > s[i+1];
    for (int i = 1; i < n; i++) rk[i] = t[i-1] && !t[i] ? (p[n1] = i, n1++) : -1;
    inducedSort(p);
    for (int i = 0, x, y; i < n; i++) if (~(x = rk[sa[i]])) {
        if (ch < 1 || p[x+1] - p[x] != p[y+1] - p[y]) ch++;
        else for (int j = p[x], k = p[y]; j <= p[x+1]; j++, k++)
            if ((s[j]<<1|t[j]) != (s[k]<<1|t[k])) {ch++; break;}
        s1[y = x] = ch;
    }
    if (ch+1 < n1) sais(n1, ch+1, s1, t+n, p+n1);
    else for (int i = 0; i < n1; i++) sa[s1[i]] = i;
    for (int i = 0; i < n1; i++) s1[i] = p[sa[i]];
    inducedSort(s1);
}
template<typename T>
int mapCharToInt(int n, const T *str) {
    int m = *max_element(str, str+n);
    fill_n(rk, m+1, 0);
    for (int i = 0; i < n; i++) rk[str[i]] = 1;
    for (int i = 0; i < m; i++) rk[i+1] += rk[i];
    for (int i = 0; i < n; i++) s[i] = rk[str[i]] - 1;
    return rk[m];
}
// Ensure that str[n] is the unique lexicographically smallest character in str.
template<typename T>
void suffixArray(int n, const T *str) {
    int m = mapCharToInt(++n, str);
    sais(n, m, s, t, p);
    for (int i = 0; i < n; i++) rk[sa[i]] = i;
    for (int i = 0, h = ht[0] = 0; i < n-1; i++) {
        int j = sa[rk[i]-1];
        while (i+h < n && j+h < n && s[i+h] == s[j+h]) h++;
        if (ht[rk[i]] = h) h--;
    }
    for(int i = 1; i < n; i++) rmq[i][0] = ht[i];
    for(int j = 1; j <= Log[n - 1]; j++) {
        for(int i = 1; i + (1 << j) <= n; i++) {
            rmq[i][j] = min(rmq[i][j - 1], rmq[i + (1 << j - 1)][j - 1]);
        }
    }
}
inline int getLcp(int L, int R) {
    assert(L < R);
    L++;
    int k = Log[R - L + 1];
    return min(rmq[L][k], rmq[R - (1 << k) + 1][k]);
}
inline pair<int, int> getLR(int p, int n, int len) {
    pair<int, int> LR(p, p);
    int low, high, mid;
    low = 1, high = p - 1;
    while(low <= high) {
        mid = low + high >> 1;
        if(getLcp(mid, p) >= len) LR.first = mid, high = mid - 1;
        else low = mid + 1;
    }
    low = p + 1, high = n;
    while(low <= high) {
        mid = low + high >> 1;
        if(getLcp(p, mid) >= len) LR.second = mid, low = mid + 1;
        else high = mid - 1;
    }
    return LR;
}
} S[2];

struct Bit {
    int a[N];
    inline void modify(int x, int v) {
        for(int i = x; i < N; i += i & -i) {
            a[i] += v;
        }
    }
    inline int sum(int x) {
        int ans = 0;
        for(int i = x; i; i -= i & -i) {
            ans += a[i];
        }
        return ans;
    }
    inline int query(int L, int R) {
        return sum(R) - sum(L - 1);
    }
} bit;

int sa_len[2];
int n, m, q, ans[N];
int pos[2][507][507];
char Map[507][507];
char t[2][N];
char str[N];
vector<int> P[2][N];

struct Qus {
    int l, r, op, id;
};
vector<int> Y[N];
vector<Qus> Q[N];

int main() {
    for(int i = 2; i < N; i++) Log[i] = Log[i >> 1] + 1;
    scanf("%d%d%d", &n, &m, &q);
    for(int i = 0; i < n; i++) scanf("%s", Map[i]);
    for(int i = 0; i < n; i++) {
        for(int j = m - 1; j >= 0; j--) {
            pos[0][i][j] = sa_len[0];
            t[0][sa_len[0]++] = Map[i][j];
        }
        t[0][sa_len[0]++] = '$';
    }
    for(int j = 0; j < m; j++) {
        for(int i = 0; i < n; i++) {
            pos[1][i][j] = sa_len[1];
            t[1][sa_len[1]++] = Map[i][j];
        }
        t[1][sa_len[1]++] = '$';
    }
    for(int i = 1; i <= q; i++) {
        scanf("%s", str);
        int len = strlen(str);
        P[0][i].resize(len);
        P[1][i].resize(len);
        for(int j = len - 1; j >= 0; j--) {
            P[0][i][j] = sa_len[0];
            t[0][sa_len[0]++] = str[j];
        }
        for(int j = 0; j < len; j++) {
            P[1][i][j] = sa_len[1];
            t[1][sa_len[1]++] = str[j];
        }
        t[0][sa_len[0]++] = '$';
        t[1][sa_len[1]++] = '$';
    }
    t[0][sa_len[0]] = '\0';
    t[1][sa_len[1]] = '\0';
    S[0].suffixArray(sa_len[0], t[0]);
    S[1].suffixArray(sa_len[1], t[1]);
    for(int i = 1; i <= q; i++) {
        for(int j = 0; j < (int)P[0][i].size(); j++) {
            pair<int, int> xLR = S[0].getLR(S[0].rk[P[0][i][j]], sa_len[0], j + 1);
            pair<int, int> yLR = S[1].getLR(S[1].rk[P[1][i][j]], sa_len[1], (int)P[0][i].size() - j);
            Q[xLR.first - 1].push_back(Qus{yLR.first, yLR.second, -1, i});
            Q[xLR.second].push_back(Qus{yLR.first, yLR.second, 1, i});
        }
    }
    for(int i = 0; i < n; i++) {
        for(int j = 0; j < m; j++) {
            int x = S[0].rk[pos[0][i][j]];
            int y = S[1].rk[pos[1][i][j]];
            Y[x].push_back(y);
        }
    }
    for(int i = 1; i <= sa_len[0]; i++) {
        for(auto &y : Y[i]) bit.modify(y, 1);
        for(auto &q : Q[i]) ans[q.id] += q.op * bit.query(q.l, q.r);
    }
    for(int i = 1; i <= q; i++) printf("%d\n", ans[i]);
    return 0;
}

/**
**/