PROBLEM LINK:
Author: yashu2040
Tester: pulkit_0110
Editorialist: freemancs
DIFFICULTY:
MEDIUM-HARD.
PREREQUISITES:
DP, Math
PROBLEM:
We want to distribute the beds to all N states such that each state gets atmost B beds.
EXPLANATION:
This problem can be solved using Dynamic Programming.
dp[i]—>>the number of ways such that we allocated i beds so far.
For each state we first iterate over number of beds used so far and then we consider all possibilites of giving beds to it.
Code for above logic
Setter's Solution
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define int long long
#define pb push_back
#define mod 1000000007
#define fast() ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL)
main()
{
// fast();
int n, B;
cin>>B>>n;
vector dp(B+1);
// dp[i] - the number of ways such that we allocated i beds so far.
dp[0]=1;
for(int states=0;states<n;states++)
{
int req;
cin>>req;
//For each state,we consider the number of beds allocated so far(used).
for(int used=B;used>=0;used–)
{
int tmp = dp[used];
//Each state can take minimum 0 beds and maximum req beds
int L = used+1;//We allocate 1 bed
int R = used+min(req,B-used);//Or we allocate the number of beds equal to req taking into account beds remaining
for(int i=L;i<=R;i++)
{
dp[i]=(dp[i]%mod+tmp%mod)%mod;
}
}
}
//final answer is dp[k].
cout<<dp[B];
}
TIME COMPLEXITY- O(NBB) WHICH WILL THROW TLE.
How to optimise it??
Lets try to prevent the use of innermost loop.
We see that we are adding tmp to range L to R for each state.
Here,we can use prefix sums for this.We can make pref array of size B for each state and we increase the count of pref[L] by 1 and decrease the count of pref[R+1]
by 1.
After iterating over all possibilites for each state,we can simply do prefix sum of our pref array and we increase dp[i] by pref[i].
Time Complexity-O(N*B)
Setter's Solution
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define int long long
#define pb push_back
#define mod 1000000007
#define fast() ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL)
main()
{
// fast();
int n, B;
cin>>B>>n;
vector<int> dp(B+1);
// dp[i] - the number of ways such that we allocated i beds so far.
dp[0]=1;
for(int states=0;states<n;states++)
{
int req;
cin>>req;
vector<int> pref(B+1,0);
//For each state,we consider the number of beds allocated so far(used).
for(int used=B;used>=0;used--)
{
int tmp = dp[used];
//Each state can take minimum 0 beds and maximum req beds
int L = used+1;//We allocate 1 bed
int R = used+min(req,B-used);//Or we allocate the number of beds equal to req taking into account beds remaining
/*for(int i=L;i<=R;i++)
{
dp[i]=(dp[i]%mod+tmp%mod)%mod;
}*/
if(L <= R)
{
pref[L]=(pref[L]%mod+tmp%mod)%mod;
if(R + 1 <= B)
{
pref[R+1]=pref[R+1]-tmp;
if(pref[R+1]<0)
pref[R+1]+=mod;
}
}
}
int sum = 0;
for(int i=0;i<=B;i++)
{
sum=(sum%mod+pref[i]%mod)%mod;
dp[i]=(dp[i]%mod+sum%mod)%mod;
}
}
//final answer is dp[B].
cout<<dp[B]%mod;
}