PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
None
PROBLEM:
Array P is said to match Q if Q can be transformed into P using the following operation:
- Choose an index i (1 \leq i \leq |Q|) and an integer X such that X \neq Q_i. Then, insert X after the i-th element of Q.
You’re given arrays A and B. Find the number of subarrays of A that match B.
EXPLANATION:
First, we need to derive a nice enough condition that tells us when array P matches array Q.
First, we make a couple of simple observations about the process:
- Our operation only allows us to insert elements into Q. So, the original elements of Q will always remain in it, in the same order.
So, Q must be a subsequence of P. - Note that our operation doesn’t allow us to change the first element of Q: we can only insert elements after an existing element, after all.
So, P_1 = Q_1 must hold.
These conditions are necessary, but not sufficient. For example, [1, 2, 3] cannot be turned into [1, 1, 2, 3] because we can’t insert a 1 at the second position.
As such, the above example can be generalized.
Suppose the first K_1 elements of Q are all equal to Q_1, and the first K_2 elements of P are equal to Q_1 (we already know K_1, K_2 \geq 1).
Then, if K_2 \gt K_1, it’s impossible for P to match Q - we cannot insert more copies of Q_1 into the appropriate indices because of having to choose a different element.
What about if K_2 \leq K_1?
It turns out that in this case, P always matches Q.
Proof
Let’s perform the operations in reverse - that is, start with P and delete elements to obtain Q.
Note that an element P_i can be deleted only if P_i \neq P_{i-1}.We know that Q appears as a subsequence of P. Greedily choose indices in P that form Q as a subsequence: say the chosen indices are i_1, i_2, \ldots, i_M.
Now, we modify this subsequence slightly - if there’s an index i_k chosen such that i_k + 1 is not chosen and P_{i_k} = P_{i_k + 1}, replace i_k with i_k + 1 instead. Note that the chosen indices still form a subsequence that equals Q.
Repeatedly do this till it’s impossible to continue.Now, we delete elements.
Going from left to right, if P_i is not included in the subsequence, delete it.
Note that our earlier adjustment of indices ensured that the first time we delete something, it will definitely not be equal to its previous element (and so the deletion is valid).
The only time this greedy deletion causes issues, is if we delete index i, making indices i-1 and i+1 adjacent - but index i+1 is not to be deleted and yet P_{i-1} = P_{i+1}.
This is easy to avoid though: if we’re ever about to perform such a deletion, just delete i+1 instead.
This process allows us to convert P into Q via deletions, so reversing them gives us the needed sequence of insertions.
We now have a set of conditions for when P matches Q:
- P_1 = Q_1
- Q should be a subsequence of P.
- The number of contiguous occurrences of P_1 at the start of P should be not more than the number of contiguous occurrences of Q_1 at the start of Q.
These conditions can be used to count the number of subarrays of A that match the given array B.
Let’s fix L, and try to find all R such that A[L\ldots R] matches B.
To do this, note that if A[L\ldots R] matches B, then so will A[L\ldots R+1]: after all, all the conditions we have depend only on prefixes, and once they’re true cannot become false (with one notable exception which will be discussed later).
So, with L fixed, let’s try to find the minimum valid R such that A[L\ldots R] matches B.
First, If A_L \neq B_1, no valid R exists so ignore this index.
Next, B needs to appear as a subsequence of the subarray.
To find the first time this happens, simply perform the greedy subsequence matching algorithm. That is:
- Let p = L initially.
- For each j from 1 to M, find the smallest index \geq p that contains the element B_j. Let this be q.
- Set p = q+1 and repeat for the next j.
- If no valid index q exists, stop - Q doesn’t appear as a subsequence in this suffix.
Now, suppose Q does appear as a subsequence, starting from L.
Let R_0 be the first time this happens (from the above algorithm, R_0 will be the final value of p, minus one).
Let K_1 be the number of contiguous appearances of B_1 at the beginning of B, and K_2 be the number of contiguous appearances of A_L starting at index L (note that this can be precomputed for all indices of A in linear time quite easily, by iterating in reverse).
Then,
- If K_2 \leq K_1, then every R \geq R_0 is a valid endpoint, so add N-R_0 + 1 to the answer.
- If K_2 \gt K_1, no index R is valid.
- There is, however, a single exception to this. If K_1 = M (so all the elements of B are equal), then the subarray of length M starting at index L is valid, though larger ones are still invalid.
So, in this case alone, add 1 to the answer.
- There is, however, a single exception to this. If K_1 = M (so all the elements of B are equal), then the subarray of length M starting at index L is valid, though larger ones are still invalid.
The complexity of this is dominated by performing the subsequence check for every starting index, giving us \mathcal{O}(NM) time.
Since M \leq 100 this is fast enough.
TIME COMPLEXITY:
\mathcal{O}(NM) per testcase.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18
mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());
void Solve()
{
int n, m; cin >> n >> m;
vector <int> a(n + 1);
for (int i = 1; i <= n; i++){
cin >> a[i];
}
vector <int> b(m + 1);
for (int i = 1; i <= m; i++){
cin >> b[i];
}
// b must be subsequence of a[l, r]
// prefix of length x must be all a[1]
int prefix = 1;
while (prefix + 1 <= m && b[prefix + 1] == b[1]){
prefix += 1;
}
// find how many b[1]'s you have
vector <int> have(n + 2, 0);
for (int i = n; i >= 1; i--){
if (a[i] == b[1]){
have[i] = have[i + 1] + 1;
} else {
have[i] = 0;
}
}
vector <int> dp(m + 1, n + 1); // minimum index where you can match b[i, m]
int ans = 0;
// there is one special edge case
// when you have equal upto prefix
// and then the next element also equal
// you cant insert
// can insert if there was some gap before though?
// so only bad case is if 1 ... 1 and you are allowed
for (int i = n; i >= 1; i--){
vector <int> ndp(m + 1, n + 1);
for (int j = 1; j <= m; j++){
if (a[i] == b[j]){
if (j == m){
ndp[j] = i;
} else {
ndp[j] = dp[j + 1];
}
} else {
ndp[j] = dp[j];
}
}
dp = ndp;
if (a[i] == b[1] && have[i] <= prefix){
ans += n + 1 - dp[1];
} else if (a[i] == b[1] && have[i] > prefix && prefix == m){
// allowed one additional array?
ans += 1;
}
}
cout << ans << "\n";
}
int32_t main()
{
auto begin = std::chrono::high_resolution_clock::now();
ios_base::sync_with_stdio(0);
cin.tie(0);
int t = 1;
// freopen("in", "r", stdin);
// freopen("out", "w", stdout);
cin >> t;
for(int i = 1; i <= t; i++)
{
//cout << "Case #" << i << ": ";
Solve();
}
auto end = std::chrono::high_resolution_clock::now();
auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n";
return 0;
}
Tester's code (C++)
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)
template<typename T>
void amin(T &a, T b) {
a = min(a,b);
}
template<typename T>
void amax(T &a, T b) {
a = max(a,b);
}
#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif
/*
*/
const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;
void solve(int test_case){
ll n,m; cin >> n >> m;
vector<ll> a(n+5), b(m+5);
rep1(i,n) cin >> a[i];
rep1(i,m) cin >> b[i];
vector<pll> blocks;
blocks.pb({0,0});
rep1(i,m){
if(b[i] == b[i-1]){
blocks.back().ss++;
}
else{
blocks.pb({b[i],1});
}
}
ll siz = sz(blocks)-1;
if(siz == 1){
ll ans = 0;
ll c = blocks[1].ss;
deque<ll> dq;
rev(i,n,1){
if(a[i] == blocks[1].ff){
dq.push_front(i);
if(sz(dq) >= c){
if(sz(dq) >= c+1 and dq[c] == i+c){
ans++;
}
else{
ans += n-dq[c-1]+1;
}
}
}
}
cout << ans << endl;
return;
}
ll dp[n+5][siz+5];
memset(dp,0x3f,sizeof dp);
deque<ll> dq[m+5];
rep(i,n+5) dp[i][siz+1] = i-1;
rev(i,n,1){
dq[a[i]].push_front(i);
for(int j = 2; j <= siz; ++j){
auto [v,c] = blocks[j];
if(sz(dq[v]) >= c){
dp[i][j] = dp[dq[v][c-1]+1][j+1];
}
}
// j = 1
if(a[i] == blocks[1].ff){
auto [v,c] = blocks[1];
if(sz(dq[v]) >= c){
if(sz(dq[v]) >= c+1 and dq[v][c] == i+c){
}
else{
dp[i][1] = dp[dq[v][c-1]+1][2];
}
}
}
}
ll ans = 0;
rep1(i,n){
ll add = max(n-dp[i][1]+1,0ll);
ans += add;
}
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
return 0;
}
Editorialist's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n, m; cin >> n >> m;
vector a(n, 0), b(m, 0);
for (int &x : a) cin >> x;
for (int &x : b) cin >> x;
vector cont(n, 1);
for (int i = n-2; i >= 0; --i) {
if (a[i] == a[i+1]) cont[i] += cont[i+1];
}
int pref = 0;
while (pref < m and b[pref] == b[0]) ++pref;
vector<array<int, 101>> jump(n);
for (int i = n-1; i >= 0; --i) {
if (i < n-1) jump[i] = jump[i+1];
else ranges::fill(jump[i], n);
jump[i][a[i]] = i;
}
ll ans = 0;
for (int i = 0; i < n; ++i) {
if (a[i] != b[0]) continue;
if (cont[i] > pref) {
ans += pref == m;
continue;
}
int p = i;
for (int j = 0; j < m; ++j) {
if (p == n or jump[p][b[j]] == n) {
p = -1;
break;
}
else p = jump[p][b[j]] + 1;
}
if (p != -1) ans += n - p + 1;
}
cout << ans << '\n';
}
}