PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Sorting
PROBLEM:
There are N items, in ascending order of size. Item i costs C_i.
You can display several items on the storefront, following these conditions:
- At least two items should be displayed when a customer visits.
- An item that’s displayed cannot be removed, unless a customer buys it.
- Each item can be bought by at most one customer; and if it is bought, is to be removed from the storefront and cannot be placed there again.
K customers will visit, in order. Each of them will buy either the largest item or the smallest item that’s on the storefront, represented by a binary string A.
Find the maximum possible revenue you can obtain by placing items appropriately.
EXPLANATION:
Exactly K items will be sold, so of course the ideal situation is if we manage to sell the K items with highest costs.
This, however, is not always possible - so the first step should be to figure out exactly which subsets of items can be sold at all.
Let’s start off with a simpler state of things: suppose every customer buys only the smallest available item (that is, A_i = 0 for all i).
Then, observe that no matter what, it’s impossible for us to sell the largest item, i.e. item N.
This is because of the requirement that there should always be at least two items on display: so if the largest item is on display, then something else will also be on display, and that item will be bought since it’s smaller.
On the other hand, we can always sell any K items excluding the N-th.
A simple method of doing this is as follows: suppose we want to sell items i_1, i_2, \ldots, i_K.
Then:
- Place i_1 and i_2 on display.
The first customer will buy i_1. - Then place i_3 on display, next to i_2 which is already there.
The second customer will buy i_2. - Then place item i_4 on display, so the third customer will buy i_3, and so on for K-1 rounds.
- Before the very last round, only item i_K will remain on display.
Place item N next to it, and the final customer will buy i_K, so that everything that we wanted to sell will be sold.
So, if all the A_i are 0, the answer is just the sum of the K largest costs, excluding item N.
Of course, if all the A_i are 1 instead, the opposite is true: the answer is the sum of the K largest costs excluding item 1, since we’re unable to sell the smallest item now.
Let’s now attempt to generalize this idea to an arbitrary binary string.
Without loss of generality, let A_K = 0, i.e. the last item bought is “small”.
Let c denote the number of times a small item is bought at the end of the string, i.e. the longest suffix of A whose elements are all equal to 0.
For example, if S = 01010011000 then c = 3, since the last three characters are equal to 0.
Then, it can be observed that no matter what we do, it’s impossible to buy all of the largest K-c+1 items.
The reasoning is just an extension of what we did in the earlier case (which had c = K):
- After the first K-c customers, at least one of the largest K-c+1 item will remain, obviously.
In particular, the largest remaining item will be one of these. - The last c customers will all buy “small” items, which means it’s impossible for the largest remaining item to be sold.
Whenever one obtains a condition like this, it’s a good idea to check if it goes the other way too: that is, is it possible for us to sell any subset of items that exclude at least one of the largest K-c+1?
It turns out, yes, we can!
Proof
The construction is not too hard.
Suppose we want to ensure items i_1, i_2, \ldots, i_K are sold, where i_1 \lt i_2\ldots\lt i_K
Let x be the largest unsold item.
By assumption, we have x \geq N - (K-c).Note that c = K has already been considered, so we only deal with c \lt K here.
In particular, this means the string has a 1 just before the last block of c zeros.Now, do the following:
- K - (c+1) times, just keep the largest two remaining items (that need to be sold) on display.
One of them will be bought and the other will remain.- When we’re down to c+1 items, a little care is needed.
At this point, there’s one item remaining on the shelf, let its index be y.
There are two cases:
- y \lt x.
This is the easy case, and is a direct extension of K = c.
Here, you can just keep all remaining items (that need to be sold) on the shelf.
All but one of them will then be bought (in some order), and to sell the last remaining item you can place x on the shelf alongside it.- y \gt x.
Here, we use the fact that K \gt c, and we’re at c+1 items so the next item being bought is a large one.
Place item x on the shelf. This will ensure the next customer buys y.
After this, place all remaining items that need to be bought on the shelf.
Because x \geq N-(K-c), all these new items will be \lt x, and hence they’ll all be bought by the last c customers, as required.
It’s now clear that we can sell any K items such that at least one of the last K-c+1 items is not among them, i.e. at least one of items N, N-1, N-2, \ldots, N-(K-c) is missing (given that A_K = 0, of course).
The aim is now to maximize the revenue following this.
This turns out to be fairly straightforward.
First, start off with just attempting to sell the K items with highest cost.
There are then two cases:
- Among these K items, at least one of N, N-1, \ldots, N-(K-c) is already missing.
In this case, this is of course optimal, so we’re done. - All the last K-c+1 items are present among the chosen K.
Here, we definitely need to discard at least one of them, and then include one other item.
The optimal choice is obviously to discard the one that costs the least, and then include whichever previously unchosen element had the highest cost.
Both cases are easily handled by just sorting the elements by cost (while remembering their original indices).
The solution above was assuming A_K = 0.
To handle A_K = 1 instead, note that basically everything is the same - just that we shouldn’t have all of the smallest K-c+1 items instead.
A simple way to handle this is: if A_K = 1, reverse the array C, and flip every 0 to 1 and vice versa in A. This then transforms the problem into an equivalent one with A_K = 0 instead, so just solve that one as described above. This avoids needing to reuse any code.
TIME COMPLEXITY:
\mathcal{O}(N\log N) 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, k; cin >> n >> k;
vector c(n, 0);
for (int &x : c) cin >> x;
string a; cin >> a;
if (a.back() == '1') {
ranges::reverse(c);
for (auto &x : a) x ^= '0'^'1';
}
int suf = 0;
for (int i = 0; i < k; ++i) {
if (a[k-1-i] != a[k-1]) break;
++suf;
}
vector ord(n, 0);
iota(begin(ord), end(ord), 0);
ranges::sort(ord, [&] (int i, int j) {return c[i] > c[j];});
ll ans = 0, have = 0;
int sub = 1e9;
for (int i = 0; i < k; ++i) {
ans += c[ord[i]];
if (ord[i] >= n-(k-suf)-1) {
have += 1;
sub = min(sub, c[ord[i]]);
}
}
if (have == k-suf+1) {
ans += c[ord[k]] - sub;
}
cout << ans%1000000007 << '\n';
}
}
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,k; cin >> n >> k;
vector<ll> a(n+5);
rep1(i,n) cin >> a[i];
string t; cin >> t;
t = "$" + t;
if(t[k] == '0'){
rep1(i,k){
if(t[i] == '0') t[i] = '1';
else t[i] = '0';
}
reverse(a.begin()+1,a.begin()+n+1);
}
ll s = 0;
rev(i,k,1){
if(t[i] == '1') s++;
else break;
}
vector<pll> b;
rep1(i,n) b.pb({a[i],i});
sort(rall(b));
b.insert(b.begin(),{0,0});
vector<bool> taken(n+5);
ll sum = 0;
rep1(i,k){
auto [val,ind] = b[i];
sum += val;
taken[ind] = 1;
}
ll pref_cnt = 0;
rep1(i,n){
if(!taken[i]) break;
else pref_cnt++;
}
ll ans = sum;
if(pref_cnt > k-s){
ll mn = inf2;
rep1(i,k-s+1) amin(mn,a[i]);
ans = sum-mn+b[k+1].ff;
}
ans %= MOD;
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
return 0;
}