PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Simple
PREREQUISITES:
None
PROBLEM:
A binary array B is called winning if all its elements can be made 1 by repeating the following process:
- Choose a subarray with strictly more ones than zeros, and set all its elements to 1.
Given a binary array A, count the number of its winning subarrays.
EXPLANATION:
Let’s try to ascertain when exactly a binary array B of length \gt 1 is winning.
First, suppose B has two adjacent ones, i.e, B_i = B_{i+1} = 1 for some index i.
Then, B is always winning: repeatedly choose a block of at least two ones and a single zero adjacent to the block, which will convert the zero into a 1.
That leaves us with arrays where each pair of ones has several zeros between them.
If there’s a pair of ones with exactly one 0 between them, forming the subarray [1, 0, 1], then B is still winning: operating on this subarray will turn it into [1, 1, 1], following which we have adjacent ones and we’re done.
We’re now left with only arrays such that each pair of ones have at least two zeros separating them.
All such arrays are not winning: in fact, we can’t even perform a single move, since every subarray of length \gt 1 will have at least as many zeros as ones (do you see why?).
This gives us a classification of winning arrays: an array is winning if and only if it either contains two consecutive ones, or [1, 0, 1] as a subarray.
With this in hand, let’s try to count all winning subarrays of A.
Suppose we fix R, the right endpoint of the subarray. We’ll try to find all left endpoints such that [L, R] is winning.
Note that if [L, R] with L \lt R is winning, so are [L-1, R], [L-2, R], \ldots, [1, R].
This is because, if [L, R] contains consecutive ones, so will any subarray starting before L; the same applies to containing [1, 0, 1].
So, we really only need to find the rightmost L such that [L, R] is valid.
To do this, clearly it suffices to find the rightmost occurrence of either [1, 0, 1] or [1, 1] as subarrays before R.
This can be done in several ways:
- You can store all occurrences of [1, 1] and [1, 0, 1] subarrays in a list, and find the last one \leq R by binary searching on this list.
- Alternately, if R is iterated from left to right, you can simply store the last occurrences of both subarray types ([1, 1] and [1, 0, 1]); and update them with the ones ending at R if applicable.
Finally, note that this entire discussion applies to only subarrays of length \gt 1.
A length 1 array is good if and only if its single element is 1, so in the end add the number of ones in A to the answer.
TIME COMPLEXITY:
\mathcal{O}(N) 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; cin >> n;
vector <int> a(n);
for (auto &x : a) cin >> x;
int ans = 0;
int f = n;
for (int i = n - 1; i >= 0; i--){
if (a[i] == 1){
ans++;
if (i + 1 < n && a[i + 1] == 1){
f = i + 1;
}
if (i + 2 < n && a[i + 2] == 1){
f = min(f, i + 2);
}
}
ans += (n - f);
}
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; cin >> n;
vector<ll> a(n+5);
rep1(i,n) cin >> a[i];
// find smallest pos to right s.t dis between 2 consecutive 1s <= 2 in this range
vector<ll> dp(n+5,inf2);
ll closest_one = inf2;
rev(i,n,1){
if(a[i]){
if(closest_one-i <= 2){
dp[i] = closest_one;
}
closest_one = i;
}
amin(dp[i],dp[i+1]);
}
ll ans = 0;
rep1(i,n){
ans += a[i];
ll j = dp[i];
if(j <= n) ans += n-j+1;
}
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
return 0;
}
Editorialist's code (PyPy3)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = 0
l = -1
for i in range(n):
if a[i] == 1:
if i > 0 and a[i-1] == 1: l = max(l, i-1)
if i > 1 and a[i-2] == 1: l = max(l, i-2)
ans += l+1
print(ans + a.count(1))