Suffix Automaton
(content/string/suffix_automaton.h)

• View this file on GitHub
• Last update: 2024-02-22 10:57:51-08:00

(Incomplete list of) Useful things to know about SAM:

• The suffix link tree (SLT) is a suffix trie that has its paths compressed (i.e. a path with no branches will get compressed to one edge).
• To find number of distinct substrings of multiple strings, just reset SAM.last to 0 in between strings.

TODO: add fn to construct suffix link tree

Depends on

• My starter code (content/utils/template.h)

Code

#pragma once
#include "../utils/template.h"

/**
 * @brief Suffix Automaton
 * @docs docs/suffix_automaton.md
 */

// TODO: add code to rebuild suffix link tree

struct suffix_automaton{
	struct st{
		int len, par;
		ll cnt;
		map<char, int> ch;
	};
	vector<st> v;
	int sz, last;
	
	void init(int n){
		v.resize(n*2);
		v[0].len = 0; //0 is the root
		v[0].par = -1;
		sz = 1;
		last = 0;
	}
	
	void extend(char c){
		int cur = sz++;
		v[cur].len = v[last].len+1;
		int p = last;
		while(p != -1 && !v[p].ch.count(c)){
			v[p].ch[c] = cur;
			p = v[p].par;
		}
		if(p == -1){
			v[cur].par = 0;
		}
		else{
			int o = v[p].ch[c]; //"other"
			if(v[p].len+1 == v[o].len){
				v[cur].par = o;
			}
			else{
				int clone = sz++; //clone
				v[clone].ch = v[o].ch;
				v[clone].par = v[o].par;
				v[clone].len = v[p].len+1;
				while(p != -1 && v[p].ch[c] == o){
					//redirect all these to clone
					v[p].ch[c] = clone;
					p = v[p].par;
				}
				v[o].par = v[cur].par = clone;
			}
		}
		last = cur;
	}
	
	ll getsz(int x){
		if(v[x].cnt)
			return v[x].cnt;
		for(auto i: v[x].ch)
			v[x].cnt += getsz(i.second);
		return ++v[x].cnt;
	}
	
	// k-th lexographically least substring (empty string counts as the 0-th one)
	string kth(int cur, int k){
		assert(k < v[cur].cnt);
		if(!k) return "";
		k--;
		for(auto i: v[cur].ch){
			if(k < v[i.second].cnt)
				return i.first+kth(i.second, k);
			k -= v[i.second].cnt;
		}
		abort();
	}
};

#line 1 "content/utils/template.h"
/**
 * @brief My starter code
 */

#include <bits/stdc++.h>
#define all(x) (x).begin(), (x).end()
#define makeunique(x) (x).erase(unique((x).begin(), (x).end()), (x).end());

#ifdef LOCAL
template<typename T> void pr(T a){std::cerr<<a<<std::endl;}
template<typename T, typename... Args> void pr(T a, Args... args){std::cerr<<a<<' ',pr(args...);}
#else
template<typename... Args> void pr(Args... args){}
#endif

using namespace std;
using ll = long long;
using pii = pair<int, int>;
#line 3 "content/string/suffix_automaton.h"

/**
 * @brief Suffix Automaton
 * @docs docs/suffix_automaton.md
 */

// TODO: add code to rebuild suffix link tree

struct suffix_automaton{
	struct st{
		int len, par;
		ll cnt;
		map<char, int> ch;
	};
	vector<st> v;
	int sz, last;
	
	void init(int n){
		v.resize(n*2);
		v[0].len = 0; //0 is the root
		v[0].par = -1;
		sz = 1;
		last = 0;
	}
	
	void extend(char c){
		int cur = sz++;
		v[cur].len = v[last].len+1;
		int p = last;
		while(p != -1 && !v[p].ch.count(c)){
			v[p].ch[c] = cur;
			p = v[p].par;
		}
		if(p == -1){
			v[cur].par = 0;
		}
		else{
			int o = v[p].ch[c]; //"other"
			if(v[p].len+1 == v[o].len){
				v[cur].par = o;
			}
			else{
				int clone = sz++; //clone
				v[clone].ch = v[o].ch;
				v[clone].par = v[o].par;
				v[clone].len = v[p].len+1;
				while(p != -1 && v[p].ch[c] == o){
					//redirect all these to clone
					v[p].ch[c] = clone;
					p = v[p].par;
				}
				v[o].par = v[cur].par = clone;
			}
		}
		last = cur;
	}
	
	ll getsz(int x){
		if(v[x].cnt)
			return v[x].cnt;
		for(auto i: v[x].ch)
			v[x].cnt += getsz(i.second);
		return ++v[x].cnt;
	}
	
	// k-th lexographically least substring (empty string counts as the 0-th one)
	string kth(int cur, int k){
		assert(k < v[cur].cnt);
		if(!k) return "";
		k--;
		for(auto i: v[cur].ch){
			if(k < v[i.second].cnt)
				return i.first+kth(i.second, k);
			k -= v[i.second].cnt;
		}
		abort();
	}
};

Back to top page