PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: dannyboy1204
Tester and Editorialist: iceknight1093
DIFFICULTY:
2929
PREREQUISITES:
Dynamic programming, observation
PROBLEM:
There are N shops; the i-th shop sells items with value V_i.
Each shop has an infinite supply of items, and each item costs exactly one dollar.
You can buy items worth at most K dollars. However, the following constraints must be followed:
- Suppose you but \text{num}_i items from the i-th shop.
- Then:
- \text{num}_1 = 0 or \text{num}_i = 1
- For i \gt 1, \text{num}_i \leq \text{num}_{i-1} + 1
Find the maximum total value of items you can buy.
EXPLANATION:
With bounds on the number of items you can take and an objective of maximizing values, this problem seems somewhat similar to the knapsack problem.
And indeed, this task can indeed be solved with dynamic programming, with each subtask requiring a slightly higher level of optimization.
Subtask 1
Let’s start with the easiest one.
A very natural dp state should encapsulate:
- Which shop we’re at, say i.
- The total number of items we’ve bought so far, say x.
- \text{num}_i, the number of items we’re buying from this shop.
So, let dp(i, x, y) be the maximum value we can obtain if we bought exactly x items from the first i shops, and of them, bought exactly y items from the i-th shop.
Buying y items from the i-th shop gives us a value of y\cdot V_i.
Further, we must’ve bought exactly x - y items from the previous shops; and in also at least y-1 items from shop i-1.
This gives us the transitions:
where the maximum is across all y-1 \leq z \leq K.
There are \mathcal{O}(NK^2) dp states here (i is upto N, x and y are upto K)2 and the transition from each one can be done in \mathcal{O}(K) time for an overall \mathcal{O}(NK^3) solution.
Subtask 2
Now N and K are a bit higher, so the previous solution is too slow.
Let’s look at the dp transitions.
We have
across y-1 \leq z.
The slow part is computing the maximum across all z.
Note that when i, x, y are fixed, what we’re looking for is a suffix maximum of the dp(i-1, x-y, \cdot ) table.
So, let’s just calculate these suffix maximums as well!
Define \text{suff}(i, x, y) = \max(dp(i, x, y), dp(i, x, y+1), dp(i, x, y+2), \ldots, dp(i, x, K)).
Then, \text{suff}(i, x, y) = \max(dp(i, x, y), \text{suff}(i, x, y+1)); so these suffix maximums can be computed as and when we compute the dp values.
With this in hand, we have dp(i, x, y) = y\cdot V_i + \text{suff}(i-1, x-y, y-1).
So, there are now \mathcal{O}(NK^2) states, and an \mathcal{O}(1) transition from each.
This results in a \mathcal{O}(NK^2) solution, which is good enough for this subtask.
Subtask 3
N and K are now large enough that we can’t even have NK^2 states anymore; it’ll take too much memory and generally be too slow.
Let’s try to cut it down a bit.
Suppose we buy y items from a shop.
Since we can increase the maximum by only a step each time, this means we must’ve bought 1, 2, 3, \ldots, y-1 items from other shops before this.
In particular, if we buy y items from a shop, we’ve bought a total of at least \frac{y \cdot (y+1)}{2} items in total.
We’re only allowed to buy K items, so we don’t really care about very large y.
In particular, we don’t care about any y such that \frac{y \cdot (y+1)}{2} \gt K; which limits us to about y \leq \sqrt{2K}.
In other words, we only care about dp(i, x, y) for about \mathcal{O}(NK\sqrt {K}) states; each of which has a \mathcal{O}(1) transition using the suffix trick from the previous subtask.
Implementing this optimization is enough to get AC.
Note that depending on your implementation, a solution with correct time complexity for subtasks 2 and 3 might be inexplicably much slower than other similar solutions.
This is most likely to be a memory-related issue.
You can optimize the memory requirement down from \mathcal{O}(NK^2) or \mathcal{O}(NK\sqrt{K}) down to \mathcal{O}(K^2) or \mathcal{O}(K\sqrt {K}) (or even \mathcal{O}(K) using a slightly different approach) by noting that dp(i, \ldots) depends only on dp(i-1, \ldots) — so, it’s enough to keep only the previous row of the dp table, instead of all of them.
This is (maybe) not required to get AC, though you should likely see a major speedup on using it.
In slower languages like Python, it might be necessary to get AC.
TIME COMPLEXITY
\mathcal{O}(NK\sqrt{K}) per test case.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
#define ll long long
#define int long long
#define fi first
#define se second
#define mat vector<vector<ll>>
using namespace std;
void db() {cout << endl;}
template <typename T, typename ...U> void db(T a, U ...b) {cout << a << ' ', db(b...);}
#ifdef Cloud
#define file freopen("input.txt", "r", stdin), freopen("output.txt", "w", stdout)
#else
#define file ios::sync_with_stdio(false); cin.tie(0)
#endif
auto SEED = chrono::steady_clock::now().time_since_epoch().count();
mt19937 rng(SEED);
const int N = 2e5 + 5, mod = 998244353, inf = 1ll << 30;
signed main(){
file;
int n, K;
cin >> n >> K;
int dp[n + 1][K + 1], a[n];
for (int &i : a) cin >> i;
for (auto &i : dp) for (int &j : i) j = 0;
int SQ = sqrt(K) * 2 + 5;
for (int i = n - 1; i >= 0; i--){
for (int j = 0; j <= K; j++){
dp[i][j] = dp[i + 1][j];
int sum = 0;
for (int k = i; k < min(n, i + SQ); k++){
sum += a[k];
int len = k - i + 1;
if (len > j) break;
dp[i][j] = max(dp[i][j], dp[i + 1][j - len] + sum);
}
if (j) dp[i][j] = max(dp[i][j], dp[i][j - 1]);
//db(i, j, dp[i][j]);
}
}
cout << dp[0][K] << '\n';
}
Editorialist's code (Python)
n, k = map(int, input().split())
a = list(map(int, input().split()))
lim = 50
dp = [ [-1 for _ in range(lim)] for _ in range(k+1)]
dp[0][0] = 0
for i in range(n):
for x in reversed(range(k+1)):
for y in reversed(range(1, min(i+2, x+1, lim))):
if dp[x-y][y-1] != -1:
dp[x][y] = y*a[i] + dp[x-y][y-1]
if y+1 < lim: dp[x][y] = max(dp[x][y], dp[x][y+1])
dp[x][0] = max(dp[x][0], dp[x][1])
ans = 0
for x in range(k+1):
for y in range(lim):
ans = max(ans, dp[x][y])
print(ans)
.