PROBLEM LINK:
Author: Naman Jain
Tester: Felipe Mota
DIFFICULTY:
MEDIUM-HARD
PREREQUISITES:
KMP, Heavy Light Trick, Euler Tour of a Tree
PROBLEM:
You are given two strings P and Q. Consider some uniformly randomly and independently chosen non-empty prefix X and non-empty suffix Y of P. A string Z is formed by concatenating X and Y in this order. Find the probability that Z is a substring of Q.
QUICK EXPLANATION:
With the help of KMP, the problem can be reduced to a data-structure problem on trees which can then be solved using heavy light trick.
EXPLANATION:
Let n = |P| , m = |Q|
Let x be a non-empty prefix of P ending at i. We calculate the number of possible y s.t. y is a suffix of P and x+y is a substring of Q. For this, we see all occurrences of x in Q. Let one such occurrence be ending at index j. Then we need to consider prefixes of Q[(j+1)....m] which are suffixes of P, Let this set be A.
Consider a tree T_1, of size n + m + 1, constructed as follows :
- Node 0 is the root
- For 1 \le v \le n, the parent[v] = \pi_P[v]. [ \pi denotes the prefix function, read more about it here ]
- For n+1 \le v \le n+m, parent[v] = length of longest suffix of Q[1...(v-n)], which is a prefix of P.
In other words, nodes of T_1 represent the prefixes of P and Q, the the parent of a node is its longest suffix which is a prefix of P.
Similarly, we can construct T_2, with reverse(P) and reverse(Q)
We can observe that all occurrences of x in Q are represented by the Q-nodes [nodes representing substrings of Q] in the subtree rooted at i in T1. Also, the set A described above is equivalent to the ancestors of the node (i+j+1) [ the node representing suffix of Q starting at j+1 ] in T_2.
So basically what we need to do is, for every Q-node in the subtree of x in T_1, we want to collect the ancestors of some node in T_2, and calculate the size of the set of all these ancestors.
So this reduces to a problem similar to the following :
You are given 2 trees T_1, T_2 and a mapping of nodes from T_1 to T_2. For a node u in T_1, let B(u) be the set of nodes in the subtree rooted at u in T_1. Let C(u) be the set of nodes corresponding [according to mapping] to B(u) in T_2. Let D(u) be the union of the ancestors of C(u) in T_2. Then find the cardinality of D(u) for all u.
There can be several approaches to solve this reduced problem. Here is one of them :
The set D(u) for a node u can be formed by merging the sets D(child_u) for all children child_u of u. We can optimize this merging step by using heavy light trick.
During the heavy light dfs, We maintain the set C described above, and calculate |D(u)| for all nodes. For adding a node to the set , we need to find the number of extra ancestors it would contribute to D(u). The can be efficiently done by keeping the set C(u) ordered according to the euler order on T_2. Now the contribution of adding a node to this set can be calculated by finding the lca with the neighboring nodes in the current set and choosing the closer one of them. The contribution will be the distance to this lca. This solution has a complexity of O(N \cdot log^2N), the log factors coming from the heavy-light trick and querying/maintaining the set of nodes.
Here are some resources to the algorithms, data-structures used in the solution :
The AC submissions in the contest have several interesting approaches to the problem, many of them quite different from the solution described above. Feel free to check them out.
SOLUTIONS:
Setter's Solution
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
using namespace std;
template<class T> ostream& operator<<(ostream &os, vector<T> V) {
os << "[ "; for(auto v : V) os << v << " "; return os << "]";}
template<class L, class R> ostream& operator<<(ostream &os, pair<L,R> P) {
return os << "(" << P.first << "," << P.second << ")";}
#define TRACE
#ifdef TRACE
#define trace(...) __f(#__VA_ARGS__, __VA_ARGS__)
template <typename Arg1>
void __f(const char* name, Arg1&& arg1){
cout << name << " : " << arg1 << std::endl;
}
template <typename Arg1, typename... Args>
void __f(const char* names, Arg1&& arg1, Args&&... args){
const char* comma = strchr(names + 1, ',');cout.write(names, comma - names) << " : " << arg1<<" | ";__f(comma+1, args...);
}
#else
#define trace(...) 1
#endif
#define ll long long
#define ld long double
#define vll vector<ll>
#define pll pair<ll,ll>
#define vpll vector<pll>
#define I insert
#define pb push_back
#define F first
#define S second
#define endl "\n"
#define vi vector<int>
#define pii pair<int, int>
#define vpii vector< pii >
const int mod = 998244353;
inline int mul(int a,int b){return (a*1ll*b)%mod;}
inline int add(int a,int b){a+=b;if(a>=mod)a-=mod;return a;}
inline int power(int a,int b){int rt=1;while(b>0){if(b&1)rt=mul(rt,a);a=mul(a,a);b>>=1;}return rt;}
inline int inv(int a){return power(a,mod-2);}
inline void modadd(int &a,int &b){a+=b;if(a>=mod)a-=mod;}
const int M = 2e5+15;
void Kmp(string& s, int l, vi& pref) {
pref[0] = 0;
int cur = 0;
for(int i=1; i<l; i++){
while (true) {
if(s[i] == s[cur]) { pref[i] = ++cur; break; }
if(cur == 0){ pref[i] = 0; break; }
cur = pref[cur-1];
}
}
}
int t1, t2;
vector<vi> g1(M), g2(M);
vi par1(M), par2(M);
int n, m;
void make_tree(int& t, vector<vi>& g, string& p, string& q, vi& par) {
vi pref(n);
Kmp(p, n, pref);
t = n + m + 1;
par[0] = 0;
for(int i=1; i<=n; i++) {
par[i] = pref[i-1];
g[par[i]].pb(i);
}
int cur = 0;
for(int i=n+1; i<t; i++) {
int x = i - (n+1);
while (true) {
if (cur == n) { cur = pref[cur-1]; }
if (cur >= n) assert(0);
if (q[x] == p[cur]) { par[i] = ++cur; break; }
if (cur == 0) { par[i] = 0; break; }
cur = pref[cur-1];
}
g[par[i]].pb(i);
}
}
const int L = 18; // log2(2e5) = 17.61
int anc[M][L];
int dep[M];
int euler[M];
int cnt;
void dfsT2(int x) {
euler[x] = cnt++;
for(auto z : g2[x]) {
dep[z] = dep[x] + 1;
dfsT2(z);
}
}
void processT2() {
for (int j=0; j<L; j++) {
for (int i=0; i<t2; i++) {
if (j==0) anc[i][0] = par2[i];
else anc[i][j] = anc[anc[i][j-1]][j-1];
}
}
dep[0] = 0;
cnt = 0;
dfsT2(0);
}
int lcaT2(int x, int y) {
if(dep[x] < dep[y]) swap(x, y);
int dif = dep[x] - dep[y];
int a;
for(int i=0, a=1; dif>0; i++, a<<=1) {
if(dif & a) {
x = anc[x][i]; dif ^= a;
}
}
if(x == y) return x;
for(int i=L-1; i>=0; i--) {
if(anc[x][i] != anc[y][i]) {
x = anc[x][i];
y = anc[y][i];
}
}
return anc[x][0];
}
int sub[M], hc[M];
void dfsT1(int x) {
sub[x] = 1;
pii mx = {0, -1};
for(auto z : g1[x]) {
dfsT1(z);
sub[x] += sub[z];
mx = max(mx, {sub[z], z});
}
hc[x] = mx.S;
}
void processT1(){
dfsT1(0);
}
int cur_cnt = 0;
set<pii> nodes;
inline int convert(int x) {
x -= n; x = m + 1 - x;
x += n; return x;
}
void add(int x) {
x = convert(x);
int ind = euler[x];
set<pii>::iterator nr;
int lca = -1;
if( (nr = nodes.lower_bound({ind, 0})) != nodes.end() ) {
lca = lcaT2(nr->S, x);
}
if ( nr != nodes.begin() ) {
--nr;
int tmp = lcaT2(nr->S, x);
if((lca==-1) || dep[tmp] > dep[lca]) lca = tmp;
}
int ext = (dep[x]-1) - ((lca == -1) ? -1: dep[lca]) ;
cur_cnt += ext;
nodes.I({ind, x});
}
void dfsAdd(int x) {
for (auto z:g1[x]) {
dfsAdd(z);
}
if(x > n && x < n+m) add(x+1);
}
int Ans;
void dfsHL(int x) {
for (auto z : g1[x]) {
if (z == hc[x]) continue;
dfsHL(z);
cur_cnt = 0;
nodes.clear();
}
if(hc[x] !=-1) dfsHL(hc[x]);
for(auto z : g1[x]) {
if (z == hc[x]) continue;
dfsAdd(z);
}
if (x > n && x < n + m) add(x+1);
if (x !=0 && x<=n) {
if(cur_cnt > 0) Ans = add(Ans, cur_cnt - 1);
}
}
void solve() {
string P, Q;
cin >> P >> Q ;
string revP = P; reverse(revP.begin(), revP.end());
string revQ = Q; reverse(revQ.begin(), revQ.end());
n = P.size();
m = Q.size();
for(int i=0; i<= n+m+5; i++) {
g1[i].clear();
g2[i].clear();
}
make_tree(t1, g1, P, Q, par1);
make_tree(t2, g2, revP, revQ, par2);
processT2();
processT1();
cur_cnt = 0; nodes.clear(); Ans = 0;
dfsHL(0);
// cout<<"Count: "<< Ans<<"\n";
Ans = mul(Ans, inv(mul(n, n)));
cout<<Ans<<"\n";
}
int main(){
ios_base::sync_with_stdio(false);cin.tie(0);cout.tie(0);cout<<setprecision(25);
int T; cin>>T;
while(T--){
solve();
}
}