snippets/suffix_array.cpp at main · AJ148-dev/snippets · GitHub

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
void radix_sort(vector<pair<pair<int,int>,int> > &a){
	int n=a.size();
	{
		vi cnt(n,0);
		for(int i=0;i<n;i++)cnt[a[i].ff.ss]++;
		vi pos(n,0);
		pos[0]=0;
		for(int i=1;i<n;i++) pos[i]=pos[i-1]+cnt[i-1];
		vector<pair<pair<int,int>,int> > new_a(n);
		for(int i=0;i<n;i++){
			new_a[pos[a[i].ff.ss]]=a[i];
			pos[a[i].ff.ss]++;
		}
		a=new_a;
	}
	{
		vi cnt(n,0);
		for(int i=0;i<n;i++)cnt[a[i].ff.ff]++;
		vi pos(n,0);
		pos[0]=0;
		for(int i=1;i<n;i++)pos[i]=pos[i-1]+cnt[i-1];
		vector<pair<pair<int,int>,int> > new_a(n);
		for(int i=0;i<n;i++){
			new_a[pos[a[i].ff.ff]]=a[i];
			pos[a[i].ff.ff]++;
		}
		a=new_a;
	}
}


//suffix_array
pair<vi,vi> suffix_array(string s){
	s+="$";
	int n=s.size();
	vector<int> p(n),c(n);
	//p->suffix array : indices of the suffixes in lexographical sorted order
	//c-> class array: class of the suffixes
	{
		vector<pair<char,int> > a(n);
		for(int i=0;i<n;i++)a[i]={s[i],i};
		sort(all(a));
		for(int i=0;i<n;i++)p[i]=a[i].ss;
		c[a[0].ss]=0;
		for(int i=1;i<n;i++){
			if(a[i].ff==a[i-1].ff)c[a[i].ss]=c[a[i-1].ss];
			else c[a[i].ss]=c[a[i-1].ss]+1;
		}
	}
	int k=0;
	while((1<<k)<n){
		vector<pair<pair<int,int>,int> > a(n);
		for(int i=0;i<n;i++)a[i]={{c[i],c[(i+(1<<k))%n]},i};
		//sort(all(a));
		radix_sort(a);
		for(int i=0;i<n;i++)p[i]=a[i].ss;
		c[a[0].ss]=0;
		for(int i=1;i<n;i++){
			if(a[i].ff==a[i-1].ff)c[a[i].ss]=c[a[i-1].ss];
			else c[a[i].ss]=c[a[i-1].ss]+1;
		}
		k++;
	}

	vector<int> lcp(n+1,0);
	{
		int k=0;
		for(int i=0;i<n-1;i++){
			int x=c[i];
			if(c[i]==0) continue;
			int j=p[x-1];
			while(s[i+k]==s[j+k])k++;
			lcp[x]=k;
			k=max(0LL,k-1);
		}
	}

	return {p,lcp};             //return as required p-> if suffix_array or c-> class_array
}