CodeChef - CHSTR 后缀数组统计


题目链接

题意
给你一个长度为 $n$ 的字符串还有 $Q$ 个询问,对于每个询问给出一个 $k$ ,你要在这个字符串的所有子串中找到 $k$ 个串并且每个串都相同的方案数。

分析:

$cnt[i]$ 为出现过 $i$ 次的子串有多少个,那么答案就是 $\sum_{i=1}^{n}cnt[i]\times C(i,k)$

那么接下来的问题就是怎么计算 $cnt$ 数组了。

考虑用 $SA$,$sa$ 数组是已经把后缀排好序了,那么观察这些后缀,有相同前缀的肯定就挨在一起,所以我们按枚举后缀 $sa[i]$,然后从每个后缀的 $Height[i]$ 扫一遍。

假设现在枚举到长度 $j$,那么

1
2
3
4
if (Height[i + 1] <= j) {
cnt[1] += n - x - j + 1;
break;
}

注意我们已经对后缀排好序,那么这个时候往后的肯定只是出现一次。

否则

1
2
3
4
5
6
for (int k = i + 1; k <= n + 1; k++) {
if (Height[k] <= j) {
cnt[k - i]++;
break;
}
}

枚举排名在后面的后缀,前面说过,有相同前缀的就会挨在一起,那么当 $Height[k]$ 比 $j$ 小时,说明有相同前缀的后缀已经统计了,直接更新 $cnt$。

然后就是 $O(n^2)$ 处理 $ans[k]$ 了。

对于每个询问输出就可以了。

总的复杂度是 $O(n^2)$。

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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
#include <bits/stdc++.h>
#include <ext/rope>
using namespace __gnu_cxx;
using namespace std;
#define mst(a,b) memset(a,b,sizeof(a))
#define ALL(x) x.begin(),x.end()
#define pii pair<int, int>
#define debug(a) cout << #a": " << a << endl;
#define eularMod(a, b) a < b ? a : a % b + b
inline int lowbit(int x){ return x & -x; }
typedef long long LL;
typedef unsigned long long ULL;
const int N = 5010;
const int mod = (int) 1e9 + 7;
const int INF = 0x3f3f3f3f;
const long long LINF = 0x3f3f3f3f3f3f3f3fLL;
const double PI = acos(-1.0);
const double eps = 1e-6;

int fac[N];
int inv[N];

int q_pow(int base, int b) {
int res = 1;
for (; b; b >>= 1, base = 1LL * base * base % mod) if (b & 1)
res = 1LL * base * res % mod;
return res;
}
void init() {
fac[0] = 1;

for (int i = 1; i < N; i++)
fac[i] = 1LL * fac[i - 1] * i % mod;

inv[N - 1] = q_pow(fac[N - 1], mod - 2);

for (int i = N - 2; i >= 0; i--)
inv[i] = 1LL * inv[i + 1] * (i + 1) % mod;
}

char s[N];
int sa[N], rak[N], Height[N], tax[N], tp[N];
int n, m;

void RSort() {
for (int i = 0; i <= m; i++) tax[i] = 0;
for (int i = 1; i <= n; i++) tax[rak[tp[i]]]++;
for (int i = 1; i <= m; i++) tax[i] += tax[i - 1];
for (int i = n; i >= 1; i--) sa[tax[rak[tp[i]]]--] = tp[i];
}
int cmp(int *f, int x, int y, int w) {
return x + w <= n && y + w <= n && f[x] == f[y] && f[x + w] == f[y + w];
}
void getSA() {
for (int i = 1; i <= n; i++) rak[i] = s[i], tp[i] = i;
m = 127, RSort();
for (int w = 1, p = 1, i; p < n; w += w, m = p) {
for (p = 0, i = n - w + 1; i <= n; i++) tp[++p] = i;
for (i = 1; i <= n; i++) if (sa[i] > w) tp[++p] = sa[i] - w;
RSort(), swap(rak, tp), rak[sa[1]] = p = 1;
for (int i = 2; i <= n; i++) rak[sa[i]] = cmp(tp, sa[i], sa[i - 1], w) ? p : ++p;
}
int j, k = 0;
for (int i = 1; i <= n; Height[rak[i++]] = k)
for (k = k ? k - 1 : k, j = sa[rak[i] - 1]; s[i + k] == s[j + k]; ++k);
}

int cnt[N];
int ans[N];

inline void add(int& x, int y) {
if ((x += y) >= mod)
x -= mod;
}
inline int C(int n, int m) {
if (n < m) return 0;
return 1LL * fac[n] * inv[n - m] % mod * inv[m] % mod;
}

int main() {
#ifdef purple_bro
freopen("in.txt", "r", stdin);
// freopen("out.txt","w",stdout);
#endif
init();

int T, q;

scanf("%d", &T);

for (;T--;) {
scanf("%d%d", &n, &q);
scanf("%s", s + 1);

mst(Height, 0);

getSA();

mst(cnt, 0);
mst(ans, 0);


for (int i = 1; i <= n; i++) {
int x = sa[i];
for (int j = Height[i]; s[x + j]; j++) {
if (Height[i + 1] <= j) {
cnt[1] += n - x - j + 1;
break;
}
for (int k = i + 1; k <= n + 1; k++) {
if (Height[k] <= j) {
cnt[k - i]++;
break;
}
}
}
}

for (int k = 1; k <= n; k++) {
for (int i = k; i <= n; i++)
add(ans[k], 1LL * cnt[i] * C(i, k) % mod);
}

for (;q--;) {
int k;

scanf("%d", &k);

if (k > n)
printf("0\n");
else
printf("%d\n", ans[k]);
}
}

return 0;
}