2021杭电多校第4场 1004.Display Substring（后缀数组）

题目链接：https://acm.hdu.edu.cn/showproblem.php?pid=6988

分析

二分答案，再check答案：
所有后缀的所有前缀即所有子串，利用前缀和二分找到小于等于当前答案的子串的数量，去除重复的，即height[i]，注意要取min

代码

#include <iostream>
#include <cstring>
#include <algorithm>

using namespace std;
typedef long long ll;

const int N = 100010;

ll k;
int n, m;
char s[N];
int sa[N], x[N], y[N], c[N], rk[N], height[N];

void get_sa()
{
	for(int i=1;i<=m;i++) c[i] = 0;
    for (int i = 1; i <= n; i ++ ) c[x[i] = s[i]] ++ ;
    for (int i = 2; i <= m; i ++ ) c[i] += c[i - 1];
    for (int i = n; i; i -- ) sa[c[x[i]] -- ] = i;
    for (int k = 1; k <= n; k <<= 1)
    {
        int num = 0;
        for (int i = n - k + 1; i <= n; i ++ ) y[ ++ num] = i;
        for (int i = 1; i <= n; i ++ )
            if (sa[i] > k)
                y[ ++ num] = sa[i] - k;
        for (int i = 1; i <= m; i ++ ) c[i] = 0;
        for (int i = 1; i <= n; i ++ ) c[x[i]] ++ ;
        for (int i = 2; i <= m; i ++ ) c[i] += c[i - 1];
        for (int i = n; i; i -- ) sa[c[x[y[i]]] -- ] = y[i], y[i] = 0;
        swap(x, y);//这里用swap要注意N开太大会让代码变慢 
        x[sa[1]] = 1, num = 1;
        for (int i = 2; i <= n; i ++ )
            x[sa[i]] = (y[sa[i]] == y[sa[i - 1]] && y[sa[i] + k] == y[sa[i - 1] + k]) ? num : ++ num;
        if (num == n) break;
        m = num;
    }
}

void get_height()
{
    for (int i = 1; i <= n; i ++ ) rk[sa[i]] = i;
    for (int i = 1, k = 0; i <= n; i ++ )
    {
        if (rk[i] == 1) continue;
        if (k) k -- ;
        int j = sa[rk[i] - 1];
        while (i + k <= n && j + k <= n && s[i + k] == s[j + k]) k ++ ;
        height[rk[i]] = k;
    }
}

int w[30],sum[100005];

bool check(int now)
{
	ll res = 0;
	for(int i=1;i<=n;i++)
	{
		int pos = upper_bound(sum + 1, sum + 1 + n, now + sum[i - 1]) - sum;
		int tmp = pos - i;
		res += tmp - min(tmp, height[rk[i]]);
	}
	if(res < k) return 0;
	return 1;
}

int main()
{
	int T = 1;
	scanf("%d",&T);
	while(T--)
	{
		scanf("%d%lld",&n,&k);
		scanf("%s", s + 1);
		for(int i=0;i<26;i++) scanf("%d",&w[i]);
		for(int i=1;i<=n;i++) sum[i] = sum[i - 1] + w[s[i] - 'a'];
		m = 122;
		get_sa();
    	get_height();
    	
    	int l = 0, r = 1e7+7;
    	while(l < r)
    	{
    		int mid = (l + r) >> 1;
    		if(check(mid)) r = mid;
    		else l = mid + 1;
		}
		if(l < 1e7+7) printf("%d\n",r);
		else printf("-1\n");
	}
    return 0;
}