PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic Programming
PROBLEM:
There are N points in a line. Some of them have lighthouses.
Each lighthouse can be given a power; a lighthouse at X with power P will light up all points Y such that |X-Y| \leq P.
Find the number of distinct ways in which the points can be lit up.
EXPLANATION:
S is the string denoting which points have lighthouses, and T denotes the final status (lit/unlit) of the points.
Let’s look at some final string T, and see when it can be reached at all.
First, suppose every point is lit up. This is (almost) always achievable: as long as even a single lighthouse exists, giving any lighthouse a power of N will light everything up.
So, the only “bad” case is when there are no lighthouses at all.
Next, suppose some point is not lit up - say, i is not lit.
Then, observe that everything to the left of i is completely independent from anything to the right of i: any lighthouse to the left of i cannot light up a point that’s \gt i without also lighting up i, and vice versa.
This independence across unlit points allows us to use dynamic programming.
Define dp_i to be the number of possible configurations considering only the first i points, such that point i is unlit (the reason for this will become clear soon).
Let’s see how dp_i can be computed.
First, if S_i = 1 (i.e. there’s a lighthouse at i), it’s impossible to have point i unlit, so dp_i = 0.
We now have S_i = 0 only.
There are two possibilities now:
- Every point before i is lit, and i alone is unlit.
- There’s some other unlit point before i.
We’ll consider each of these cases separately.
Case 1: Everything before i is lit.
This is a single configuration, so we only need to check if it’s possible or not, and add 0 or 1 to dp_i appropriately.
Obviously, if there are no lighthouses before i, it’s impossible to achieve this (unless i = 1 in which case it’s possible); so we assume a lighthouse does exist.
Suppose there’s only one lighthouse, say at index j (j \lt i).
Then, the absolute best we can do is to give it a power of i - j -1: anything more would result in i being lit up.
With a power of i-j-1, every point in [j, i-1] will be lit up for sure, so only points to the left of j need to be checked.
In fact, it’s enough to only verify that point 1 is lit, since if it is, everything else will be too.
This will happen if and only if j - (i - j - 1) \leq 1, since the leftmost lit point will be exactly j - (i - j - 1).
So, simply checking this inequality (which further reduces to 2j \leq i) will tell us whether it’s possible.
Next, suppose there are several lighthouses before i.
Then, we see that it’s optimal to only consider the leftmost lighthouse: if it can’t light up point 1, nothing else can.
So, checking for whether everything before i can be lit up, but not i, is quite simple: if j is the index of the leftmost lighthouse, just check whether 2j \leq i or not.
Case 2: There’s some unlit point before i.
Let’s fix j \lt i to be the nearest unlit point to i.
Then, note that:
- All of the points j+1, j+2, \ldots, i-1 must be lit, using only lighthouses in that range.
- Points before j are independent of what is done after j: in particular, there are dp_j configurations.
This transition is exactly why we enforced the last point to be unlit in the definition of the DP.
So, if the segment [j+1, i-1] can be lit up (without lighting up i and j), then we add dp_j to dp_i; otherwise we do nothing.
All that remains is to perform this check.
The idea here is quite similar in spirit to what we did in the first case above.
There, we always considered the leftmost lighthouse, which was fine because there was no restriction on lighting up points before 1. Now though, we’re restricted on both sides.
It’s not hard to see that what should instead be done is to choose points close to the middle of the segment [j+1, i-1].
Specifically, let m = \frac{i+j}{2} be the middle of the segment. Then,
- Note that a lighthouse at a point \gt m can never light up point j+1 without lighting up i.
This means we must choose a lighthouse in [j+1, m] to be able to light j+1.
Among them, it’s best to choose the rightmost one, since everything to its left will be lit up by it anyway. - Similarly, a lighthouse that’s \lt m cannot light up i-1, meaning we need a lighthouse in [m, i-1] as well.
Here, it’s best to choose the leftmost one, for the same greedy reason as above: everything to the right will be lit up anyway.
So, let x be the rightmost lighthouse in [j+1, m] and y be the leftmost lighthouse in [m, i-1].
Then,
- If either x or y don’t exist, no solution exists.
- Otherwise, everything \leq x and \geq y will be taken care of, by giving the lighthouses powers of x - (j+1) and (i-1)-y, respectively.
- That leaves points between x and y.
Of them, the first lighthouse will cover everything till x + x - (j+1), while the second will cover everything till y - ((i-1) - y). - Check if this is point by checking if 2x - (j+1) \geq 2y - (i-1) - 1.
If x and y above are known, the check can be performed in \mathcal{O}(1) time.
Finding them quickly is fairly easy: you can do it in \mathcal{O}(1) time by storing, for each index, the closest lighthouse to its left/right for example.
Alternately, store a list of positions of the lighthouses use binary search for \mathcal{O}(\log N) time: slower but still fast enough to pass.
Since each j \lt i is now processed in constant time, that gives \mathcal{O}(N) per index, and \mathcal{O}(N^2) overall.
We now have every dp_i value in \mathcal{O}(N^2) time.
All that remains is to combine them for the final answer.
For this, let’s go back to our two cases: either every point is lit up, or some point is not lit up.
In the first case, there’s only one configuration.
As for the second, let’s fix i to be the last unlit point. Then,
- There are dp_i configurations of points till i.
- Everything after i must be lit up.
This is the same check as when we wanted an entire prefix to be lit up: just find the rightmost lighthouse and check if it’s far enough away.
If the check passes, add dp_i to the answer.
One edge case to be wary of is when there are no lighthouses at all: the answer here is 1.
TIME COMPLEXITY:
\mathcal{O}(N^2) per testcase.
CODE:
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; cin >> n;
string s; cin >> s;
vector<int> pos;
for (int i = 0; i < n; ++i) if (s[i] == '1')
pos.push_back(i);
if (pos.empty()) {
cout << 1 << '\n';
continue;
}
vector full(n, vector(n, 0)); // full[i][j] = can i have t[i] = t[j] = 0, t[i+1] = t[i+2] = ... = t[j-1] = 1?
for (int i = 0; i < n; ++i) {
if (s[i] == '1') continue;
if (i+1 < n) {
if (s[i+1] == '0') full[i][i+1] = 1;
}
for (int j = i+2; j < n; ++j) {
if (s[j] == '1') continue;
int mid = (i + j) / 2;
// find last 1 that's <= mid and first 1 that's > mid
int x = -1, y = -1;
auto it = ranges::upper_bound(pos, mid);
if (it != begin(pos) and *prev(it) > i) x = *prev(it);
if (it != end(pos) and *it < j) y = *it;
if (x == -1) continue;
int r = x + (x - i - 1), l = j;
if (y != -1) l = y - (j - 1 - y);
full[i][j] = l <= r+1;
}
}
const int mod = 998244353;
int ans = 1;
vector dp(n, 0);
for (int i = 0; i < n; ++i) {
if (s[i] == '1') continue;
// at least two zeros in prefix
for (int j = i-1; j >= 0; --j)
dp[i] = (dp[i] + 1ll*dp[j]*full[j][i]) % mod;
// one zero in prefix
if (i == 0 or (pos[0] < i and pos[0] - (i - 1 - pos[0]) <= 0)) ++dp[i];
dp[i] %= mod;
// this is last 0
if (i == n-1 or (pos.back() > i and pos.back() + (pos.back() - i - 1) >= n-1)) ans = (ans + dp[i]) % mod;
}
cout << ans << '\n';
}
}
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());
const int mod = 998244353;
struct mint{
int x;
mint (){ x = 0;}
mint (int32_t xx){ x = xx % mod; if (x < 0) x += mod;}
mint (long long xx){ x = xx % mod; if (x < 0) x += mod;}
int val(){
return x;
}
mint &operator++(){
x++;
if (x == mod) x = 0;
return *this;
}
mint &operator--(){
if (x == 0) x = mod;
x--;
return *this;
}
mint operator++(int32_t){
mint result = *this;
++*this;
return result;
}
mint operator--(int32_t){
mint result = *this;
--*this;
return result;
}
mint& operator+=(const mint &b){
x += b.x;
if (x >= mod) x -= mod;
return *this;
}
mint& operator-=(const mint &b){
x -= b.x;
if (x < 0) x += mod;
return *this;
}
mint& operator*=(const mint &b){
long long z = x;
z *= b.x;
z %= mod;
x = (int)z;
return *this;
}
mint operator+() const {
return *this;
}
mint operator-() const {
return mint() - *this;
}
mint operator/=(const mint &b){
return *this = *this * b.inv();
}
mint power(long long n) const {
mint ok = *this, r = 1;
while (n){
if (n & 1){
r *= ok;
}
ok *= ok;
n >>= 1;
}
return r;
}
mint inv() const {
return power(mod - 2);
}
friend mint operator+(const mint& a, const mint& b){ return mint(a) += b;}
friend mint operator-(const mint& a, const mint& b){ return mint(a) -= b;}
friend mint operator*(const mint& a, const mint& b){ return mint(a) *= b;}
friend mint operator/(const mint& a, const mint& b){ return mint(a) /= b;}
friend bool operator==(const mint& a, const mint& b){ return a.x == b.x;}
friend bool operator!=(const mint& a, const mint& b){ return a.x != b.x;}
mint power(mint a, long long n){
return a.power(n);
}
friend ostream &operator<<(ostream &os, const mint &m) {
os << m.x;
return os;
}
explicit operator bool() const {
return x != 0;
}
};
// Remember to check MOD
void Solve()
{
int n; cin >> n;
string s; cin >> s;
s = "0" + s + "0";
vector <int> nx(n + 2, n + 1);
vector <int> pv(n + 1, -1);
vector <int> ps(n + 1, 0);
for (int i = 1; i <= n; i++){
ps[i] = ps[i - 1] + (s[i] - '0');
if (s[i] == '1'){
pv[i] = i;
} else {
pv[i] = pv[i - 1];
}
}
for (int i = n; i >= 1; i--){
if (s[i] == '1'){
nx[i] = i;
} else {
nx[i] = nx[i + 1];
}
}
auto good = [&](int l, int r){
if (l > r){
return true;
}
if (l == 1 && r == n){
return ps[n] != 0;
}
if (l == 1){
return ps[(l + r) / 2] != 0;
}
if (r == n){
int pos = (l + r + 1) / 2;
return ps[n] != ps[pos - 1];
}
int len = (r - l + 1);
if (len % 2 == 1){
int c = (l + r) / 2;
int x = pv[c];
int y = nx[c];
int diff = (y - x);
return (diff <= len / 2);
} else {
int c1 = (l + r) / 2;
int c2 = (l + r + 1) / 2;
int x = pv[c1];
int y = nx[c2];
int diff = (y - x);
return (diff <= len / 2);
}
};
vector <mint> dp(n + 2, 0);
dp[0] = 1;
for (int i = 0; i <= n; i++) if (s[i] == '0'){
for (int j = i + 1; j <= n + 1; j++){
if (s[j] == '0' && good(i + 1, j - 1)){
dp[j] += dp[i];
}
}
}
cout << dp[n + 1] << "\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;
}