P3975 [TJOI2015]弦论（SAM DAG、parent树上dp计算不同子串数 递归输出字典序第k大子串）

[TJOI2015]弦论

题目描述

为了提高智商，ZJY 开始学习弦论。这一天，她在《String theory》中看到了这样一道问题：对于一个给定的长度为 $n$ 的字符串，求出它的第 $k$ 小子串是什么。你能帮帮她吗?

输入格式

第一行是一个仅由小写英文字母构成的字符串 $s$。

第二行为两个整数 $t$ 和 $k$，$t$ 为 $0$ 则表示不同位置的相同子串算作一个，$t$ 为 $1$ 则表示不同位置的相同子串算作多个。$k$ 的意义见题目描述。

输出格式

输出数据仅有一行,该行有一个字符串，为第 $k$ 小的子串。若子串数目不足 $k$ 个，则输出 $-1$。

样例 #1

样例输入 #1

aabc
0 3

样例输出 #1

aab

样例 #2

样例输入 #2

aabc
1 3

样例输出 #2

aa

样例 #3

样例输入 #3

aabc
1 11

样例输出 #3

-1

提示

数据范围

对于 $10\%$ 的数据，$n\le 1000$。

对于 $50\%$ 的数据，$t=0$。

对于 $100\%$ 的数据，$1\le n\le 5\times 10^5$，$0\le t\le 1$，$1\le k\le 10^9$。

题意：

给定两种询问模式，
op = 1：不同位置的相同子串算多个。
op = 0：不同位置的相同子串算一个。
要求输出原串字典序第k大的子串。

思路：

• 先考虑 op = 0，即本质不同才算不同：

对于每个点 u，维护一个 sum[u]，表示 经过 u 的本质不同子串数量。这等价于在 SAM 的 DAG 上，u 所能到达的点的数量（等价于除根之外的点点权为 1，即直接初始化 sum[u] = cnt[u] = 1，cnt[u] 表示节点 u 代表子串结束位置集合，既然不同位置相同子串算一种，那么直接置为 1，求所能到达的点权和），可以 dfs 或 按拓扑序递推 求得，这里我直接跑一遍 dfs 进行计算。

• 现在考虑 op = 1，即位置不同即不同：

对于每个点 u，维护一个 sum[u] 表示 经过 u 的位置不同子串数量。显然经过它的每个串 s 的贡献是 endpos(s) 的大小（这就是出现的 s 次数）。而根到 u 表示串的 endpos 大小就是 parent 树上 u 的子树中的前缀点数量 cnt[u]。先 dfs 求出 cnt[u] 表示 u 的子树大小，并将它作为点权，问题即转化为 op = 0 的情况。

上面两种均为先按要求初始化 cnt 和 sum 数组，之后再进行 dfs。

• 考虑查询：(假设现在的节点是 u,要查第 k 小子串）

先将 k 减去 u 的点权 cnt[u]。按字典序递增考虑 u 能直接到达的点 v，如果 k ≤ sum[c]，则输出 c 表示的字符，并递归查询 (c, k)，否则 k 减去 sum[c]
​

注意判断 -1 的情况

总结：

• SAM 中的连边只有两种，一种是 DAG 上的（ch 指针），一种是 parent 树上的（fa 指针），前者是 有向无环图，后者是单向树。

• 一般看到 SAM 会配合基数排序然后倒着维护答案，这个过程实际上模拟的是在 parent 树上的 dfs，更直观的理解就是，利用 fa 指针将 parent 树建出来，然后直接在树上维护信息即可。

• 如果想要维护 dp 的话，需要在 DAG 上跑拓扑，在 parent 树上跑树形 dp
  顺着 DAG 跑的话可以得到子串的前缀，顺着 parent 树往上跳 fa 的话可以得到后缀的后缀。

代码：

#include<bits/stdc++.h>

using namespace std;
const int N = 5e5 + 10, M = N << 1;
int ch[M][26], len[M], fa[M], np = 1, tot = 1;
bool st[M];
char s[N];
long long cnt[M], sum[M];	//cnt表示节点代表字符串的结束位置集合 sum表示从某个节点开始继续往下走能够到达的子串个数 sum由cnt按拓扑序递推而来
vector<int> g[M];
int op, k;

void extend(int c)
{
	int p = np; np = ++tot;
	len[np] = len[p] + 1, cnt[np] = 1;
	while (p && !ch[p][c]) {
		ch[p][c] = np;
		p = fa[p];
	}
	if (!p) {
		fa[np] = 1;
	}
	else {
		int q = ch[p][c];
		if (len[q] == len[np] + 1) {
			fa[np] = q;
		}
		else {
			int nq = ++tot;
			len[nq] = len[np] + 1;
			fa[nq] = fa[q], fa[q] = fa[np] = nq;
			while (p && ch[p][c] == q) {
				ch[p][c] = nq;
				p = fa[p];
			}
			memcpy(ch[nq], ch[q], sizeof ch[q]);
		}
	}
}

void dfs1(int u)
{
	for (auto son : g[u]) {
		dfs1(son);
		cnt[u] += cnt[son];	//计算 endpos
	}
	sum[u] = op ? cnt[u] : cnt[u] = 1;	//根据 op 来初始化 sum
}

void dfs2(int u)
{
	if (st[u]) return;
	st[u] = true;	//这里必须要防止重复遍历，防止超时
	for (int i = 0; i < 26; ++i) {
		if (!ch[u][i]) continue;
		dfs2(ch[u][i]);
		sum[u] += sum[ch[u][i]];	//计算经过节点 u 的本质不同或位置不同的子串数量
	}
}

void print(int u, int k)
{
	if (k <= cnt[u]) return;
	k -= cnt[u];	//先减去点权
	for (int i = 0; i < 26; ++i) {
		int c = ch[u][i];
		if (!c) continue;
		if (sum[c] < k) {	//相当于剪枝，如果当前节点某棵子树包含的子串完全包含于前 k 大的话，直接剪掉就行了
			k -= sum[c];	//显然后续要找的就不是第 k 大了，要变化一下
			continue;
		}
		printf("%c", 'a' + i);	//符合条件 输出
		print(c, k);
		return;
	}
}

signed main()
{
	scanf("%s", s);
	scanf("%d%d", &op, &k);
	for (int i = 0; s[i]; ++i) extend(s[i] - 'a');
	for (int i = 2; i <= tot; ++i) g[fa[i]].emplace_back(i);
	dfs1(1);
	sum[1] = cnt[1] = 0;	//初始是空节点，显然初始化为 0
	dfs2(1);
	if (k > sum[1]) {	//把贡献都加在 sum[1] 表示原串所有子串的数量，这里直接判断有没有第 k 大
		printf("-1\n");
	}
	else {
		print(1, k);
	}
	puts("");

	return 0;
}