poj_3415 后缀数组+单调栈

题目大意

 
对于给定的字符串A和B，求出满足要求的集合S的大小。

题目分析

poj_3415

实现(c++)

 

#define _CRT_SECURE_NO_WARNINGS
#include<stdio.h>
#include<string.h>
#include<vector>
using namespace std;
#define LETTERS 60
#define MAX_ARRAY_SIZE 200005

int gSuffixArray[MAX_ARRAY_SIZE];
int gCount[MAX_ARRAY_SIZE];
int gOrderBySecondKey[MAX_ARRAY_SIZE];
int gRank[MAX_ARRAY_SIZE];
int gFirstKeyArray[MAX_ARRAY_SIZE];
int gHeight[MAX_ARRAY_SIZE];

int gStr[MAX_ARRAY_SIZE];
int gStrLen;

bool Compare(int* arr, int a, int b, int step){
	return arr[a] == arr[b] && arr[a + step] == arr[b + step];
}

void GetStr(char* str){
	memset(gStr, 0, sizeof(gStr));
	gStrLen = strlen(str);
	for (int i = 0; i < gStrLen; i++){
		if (str[i] >= 'a'){
			gStr[i] = str[i] - 'a' + 27;
		}
		else{
			gStr[i] = str[i] - 'A' + 1;
		}
	}
	gStr[gStrLen] = 0;
	gStrLen++;
}
//求后缀数组
void GetSuffixArray(){
	int n = gStrLen;
	memset(gCount, 0, sizeof(gCount));
	for (int i = 0; i < n; i++){
		gRank[i] = gStr[i];
		gCount[gRank[i]] ++;
	}
	int m = LETTERS;
	for (int i = 1; i < m; i++){
		gCount[i] += gCount[i - 1];
	}
	for (int i = n - 1; i >= 0; i--){
		gSuffixArray[--gCount[gRank[i]]] = i;
	}

	int step = 1;
	int *rank = gRank, *order_by_second_key = gOrderBySecondKey;
	while (step < n){
		int p = 0;

		for (int i = n - step; i < n; i++){
			order_by_second_key[p++] = i;
		}
		for (int i = 0; i < n; i++){
			if (gSuffixArray[i] >= step){
				order_by_second_key[p++] = gSuffixArray[i] - step;
			}
		}
		for (int i = 0; i < n; i++){
			gFirstKeyArray[i] = rank[order_by_second_key[i]];
		}
		for (int i = 0; i < m; i++){
			gCount[i] = 0;
		}
		for (int i = 0; i < n; i++){
			gCount[gFirstKeyArray[i]] ++;
		}
		for (int i = 1; i < m; i++){
			gCount[i] += gCount[i - 1];
		}
		for (int i = n - 1; i >= 0; i--){
			gSuffixArray[--gCount[gFirstKeyArray[i]]] = order_by_second_key[i];
		}
		int* tmp = rank; rank = order_by_second_key; order_by_second_key = tmp;
		rank[gSuffixArray[0]] = p = 0;
		for (int i = 1; i < n; i++){
			if (Compare(order_by_second_key, gSuffixArray[i], gSuffixArray[i - 1], step)){
				rank[gSuffixArray[i]] = p;
			}
			else{
				rank[gSuffixArray[i]] = ++p;
			}
		}
		m = p + 1;
		step *= 2;
	}
}
//求height数组
void GetHeight(){
	int n = gStrLen;
	for (int i = 0; i < n; i++){
		gRank[gSuffixArray[i]] = i;
	}
	int k = 0, j;
	for (int i = 0; i < n; i++){
		if (k){
			k--;
		}
		j = gSuffixArray[gRank[i] - 1];
		while (j + k < n && i + k < n&& gStr[i + k] == gStr[j + k]){
			k++;
		}
		gHeight[gRank[i]] = k;
	}
}
int min(int a, int b){
	return a < b ? a : b;
}
int gStack[MAX_ARRAY_SIZE][2];
long long int Find(int k, int n){
	int end = 1;
	long long int sum = 0, tot = 0;
	int top = -1, count = 0;
	while (end < gStrLen){		
		if (gHeight[end] < k){
			count = 0;
			tot = 0;
			top = -1;
		}
		else{
			count = 0;
			if (gSuffixArray[end - 1] < n){
				tot += (gHeight[end] - k + 1);
				count++;
			}
			while (top >= 0 && gStack[top][0] >= gHeight[end]){
				tot -= gStack[top][1] * (gStack[top][0] - gHeight[end]);
				count += gStack[top][1];
				top--;
			}
			top++;
			gStack[top][0] = gHeight[end];
			gStack[top][1] = count;
			if (gSuffixArray[end] > n){
				sum += tot;
			}
		}		
		end++;
	}
	end = 1;
	tot = 0;
	count = 0;
	top = -1;
	while (end < gStrLen){
		if (gHeight[end] < k){
			tot = 0;
			count = 0;
			top = -1;
		}
		else{
			count = 0;
			if (gSuffixArray[end - 1] > n){
				count++;
				tot += (gHeight[end] - k + 1);
			}
			while (top >= 0 && gStack[top][0] >= gHeight[end]){
				tot -= (gStack[top][1])*(gStack[top][0] - gHeight[end]);
				count += gStack[top][1];
				top--;
			}
			top++;
			
			gStack[top][0] = gHeight[end];
			gStack[top][1] = count;

			if (gSuffixArray[end] < n){
				sum += tot;
			}
		}
		end ++;
	}
	return sum;
}
char str[MAX_ARRAY_SIZE];
int main(){
	int len1, len2, k;
	while (scanf("%d", &k) != EOF){
		if (k == 0){
			break;
		}
		scanf("%s", str);
		len1 = strlen(str);
		str[len1] = 'a' + 27;
		scanf("%s", str + len1 + 1);
		GetStr(str);
		GetSuffixArray();
		GetHeight();
	
		long long int sum = Find(k, len1);
		
		printf("%lld\n", sum);
	}
	return 0;
}