204E - Little Elephant and Strings - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
const int N=2e5+10;
int n,k;
string s[N];
vector<int> pos[N];
int dfn[N];
int cmp(int x,int y){return dfn[x]<dfn[y];}
struct S{
	int las,idx;
	int ch[N][26],len[N],fail[N];
	int cnt,siz[N],dep[N],f[20][N];
	vector<int> e[N];
	S(){fail[0]=-1;}
	void insert(int c,int &x){
		if(ch[las][c]){
			int p=las,q=las=ch[p][c];
			x=q;
			if(len[q]==len[p]+1) return;
			else{
				int t=las=++idx;
				x=t;
				memcpy(ch[t],ch[q],sizeof(ch[t]));
				fail[t]=fail[q];
				len[t]=len[p]+1;
				fail[q]=t;
				for(;~p&&ch[p][c]==q;p=fail[p]) ch[p][c]=t;
			}
			return;
		}
		int p=las,now=las=++idx;
		x=now;
		len[now]=len[p]+1;
		for(;~p&&!ch[p][c];p=fail[p]) ch[p][c]=now;
		if(p==-1){
			fail[now]=0;
			return;
		}
		int q=ch[p][c];
		if(len[q]==len[p]+1) fail[now]=q;
		else{
			int t=++idx;
			memcpy(ch[t],ch[q],sizeof(ch[t]));
			fail[t]=fail[q];
			len[t]=len[p]+1;
			fail[now]=fail[q]=t;
			for(;~p&&ch[p][c]==q;p=fail[p]) ch[p][c]=t;
		}
	}
	void dfs(int x){
		dfn[x]=++cnt;
		for(auto y:e[x]){
			dep[y]=dep[x]+1;
			f[0][y]=x;
			dfs(y);
		}
	}
	void pre(){
		for(int i=1;i<=18;i++){
			for(int j=1;j<=idx;j++){
				f[i][j]=f[i-1][f[i-1][j]];
			}
		}
	}
	int lca(int x,int y){
		if(dep[x]<dep[y]) swap(x,y);
		for(int i=18;i>=0;i--){
			if(dep[f[i][x]]>=dep[y]) x=f[i][x];
		}
		if(x==y) return x;
		for(int i=18;i>=0;i--){
			if(f[i][x]!=f[i][y]){
				x=f[i][x];
				y=f[i][y];
			}
		}
		return f[0][x];
	}
	void calc(int x){
		for(auto y:e[x]){
			calc(y);
			siz[x]+=siz[y];
		}
	}
	void solve(){
		for(int i=1;i<=idx;i++) e[fail[i]].push_back(i);
		dfs(0);
		pre();
		for(int i=1;i<=n;i++){
			sort(pos[i].begin(),pos[i].end(),cmp);
			for(int j=0;j<(int)s[i].length();j++) siz[pos[i][j]]++;
			for(int j=1;j<(int)s[i].length();j++) siz[lca(pos[i][j-1],pos[i][j])]--;
		}
		calc(0);
		for(int i=1;i<=n;i++){
			long long res=0;
			for(int j=0;j<(int)s[i].length();j++){
				int p=pos[i][j];
				for(int t=18;t>=0;t--){
					if(siz[f[t][p]]<k) p=f[t][p];
				}
				if(siz[p]<k) p=f[0][p];
				if(p) res+=len[p];
			}
			cout<<res<<' ';
		}
	}
}Sam;
int main(){
	cin>>n>>k;
	for(int i=1;i<=n;i++){
		cin>>s[i];
		Sam.las=0;
		pos[i].resize(s[i].length());
		for(int j=0;j<(int)s[i].length();j++){
			Sam.insert(s[i][j]-'a',pos[i][j]);
		}
	}
	Sam.solve();
	return 0;
}