WNTR - Editorial

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: Yaswanth Phani Kommineni
Testers: Nishank Suresh, Tejas Pandey
Editorialist: Nishank Suresh

DIFFICULTY:

3120

PREREQUISITES:

Dynamic programming

PROBLEM:

You have two arrays A and B of length N. For each K from 0 to N-1, find the following:

  • Let C be an array with 0 \leq C_i \leq B_i
  • At least K elements of C must be 0
  • sum(A) = sum(C)
  • Under these constraints, find the minimum value of \sum_{i=1}^N |A_i - C_i|.

EXPLANATION:

Let’s fix a value of K and see what happens to C.

There are at most N-K non-zero indices; let S denote the set of these non-zero indices. Then, we have the following:

  • Let B_S = \sum_{i \in S}B_i. We must have B_S \geq sum(A), otherwise it’s impossible for sum(C) = sum(A) to hold. For now, let’s assume this condition holds.
  • We can assign C using the following strategy:
    • If i \notin S, assign C_i := 0.
    • If i \in S and B_i \leq A_i, assign C_i := B_i. This is the best we can do at this index. Such indices contribute a value of A_i - B_i to the result.
    • If i \in S and B_i \gt A_i, assign C_i := A_i initially.
    • Now, we need to add some values to indices of the third type to ensure sum(C) = sum(A). Note that it doesn’t really matter which index a value is added to: adding x to any such index increases the result by x.
    • Since B_S \geq sum(A), it’s always possible to achieve sum(C) = sum(A), so we only need to find out how much is added — that will tell us the result.
  • Finding out how much is added is simple: we have to add A_i for each i \notin S, and then A_i - B_i for each i\in S with A_i \geq B_i.

So, for the S we chose, the result can be computed as follows:

  • Let B_S be as denoted above. Similarly, let A_S = \sum_{i\in S}A_i.
  • Further, let D = \sum_{i\in S} \max(0, A_i - B_i).
  • Then,
    • The total contribution of indices i \in S with B_i \leq A_i is D
    • The total contribution of indices i \in S with B_i \gt A_i is sum(A) - A_S + D
    • The total contribution of indices i \notin S is sum(A) - A_S
  • Adding all these together, the result for this set S is 2\cdot(D + sum(A) - A_S)

We would like to choose S in such a way that this is minimized. sum(A) is a constant, so this is equivalent to minimizing D - A_S.
Note that a given index contributes (\max(0, A_i - B_i) - A_i) to D-A_S if it is chosen in S.

We also need to keep track of B_S, since we need to ensure it is greater than sum(A).

This allows us to write the following knapsack-style dynamic programming solution:

  • Let dp_{i, k, x} denote the minimum possible value of D-A_S for a set of size k chosen from indices 1, 2, \ldots, i, whose sum of B_i values is x.
  • At position i, we have two choices: either take index i into S, or don’t
    • If i is not taken, we simply have dp_{i-1, k, x}
    • If i is taken, we have dp_{i-1, k-1, x - B_i} + \max(0, A_i - B_i) - A_i
    • dp_{i, k, x} is the minimum of these two values.

This gives us a solution in \mathcal{O}(N^4) with a pretty low constant factor (more specifically, it is \mathcal{O}(N^2 \cdot sum(B)), but sum(B) \leq 10^4 anyway).

However, it currently uses \mathcal{O}(N^4) memory, which might be too much. Note that by iterating from left to right, when we are at i, we only care about the dp values at i-1. This allows us to cut down the memory to \mathcal{O}(N^3).

TIME COMPLEXITY

\mathcal{O}(N^4) per test case.

CODE:

Setter's code (C++)
/*
Yaswanth Phani Kommineni
*/
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;

void solve(){
        int n,s = 0;
        cin >> n;
        vector <int> a(n),b(n);
        for(int i=0;i<n;i++){
                cin >> a[i];
                s += a[i];
        }
        for(int j=0;j<n;j++){
                cin >> b[j];
        }
        vector < vector < int> > dp(n+1, vector <int> (s+1,-1e9));
        dp[0][0] = 0;
        for(int i=0;i<n;i++){
                for(int j=n;j>=1;j--){
                        for(int k=s;k>=min(a[i],b[i]);k--){
                                dp[j][k] = max(dp[j][k],dp[j-1][k-min(a[i],b[i])] + max(0,b[i]-a[i]));
                        }
                }
        }
        for(int i=n;i>=1;i--){
                int cur;
                bool fl = false;
                for(int j=0;j<=s;j++){
                        if(s-j <= dp[i][j]){
                                if(!fl){
                                        fl = true;
                                        cur = 2*(s-j);
                                        continue;
                                }
                                cur = min(cur, 2*(s-j));
                        }
                }
                if(!fl) cur = -1;
                cout << cur << " ";
        }
        cout << endl;
}

int main(){
#ifndef ONLINE_JUDGE
        //freopen("input0.in","r",stdin);
        //freopen("input0.out","w",stdout);
#endif
        ios_base::sync_with_stdio(false);
        cin.tie(0);
        cout.tie(0);
        int tc;
        tc = 1;
        cin >> tc;
        for(int i=1;i<=tc;i++){
                //cout << "Case #" << i << ": ";
                solve();
        }
        return 0;
}
Tester's code (C++)
#include <bits/stdc++.h>
using namespace std;

int main() {
	int t;
	cin >> t;
	while(t--) {
	    int n;
	    cin >> n;
	    vector<int> a(n + 1), b(n + 1);
	    int sma = 0;
	    for(int i = 1; i <= n; i++) cin >> a[i], sma += a[i];
	    for(int i = 1; i <= n; i++) cin >> b[i];
	    int dp[n + 1][sma + 1];
	    memset(dp, -1, sizeof(dp));
	    dp[0][0] = 0;
	    for(int i = 1; i <= n; i++) {
	        for(int j = n; j ; j--) {
	            for(int k = min(a[i], b[i]); k <= sma; k++) {
	                if(dp[j - 1][k - min(a[i], b[i])] > -1)
	                dp[j][k] = max(dp[j][k], dp[j - 1][k - min(a[i], b[i])] + max(0, b[i] - a[i]));
	            }
	        }
	    }
	    for(int i = n; i; i--) {
	        int ans = -1;
	        for(int j = 0; j <= sma; j++) {
	            if(dp[i][j] > -1 && sma - j <= dp[i][j] && (ans == -1 || ans > 2*(sma - j))) ans = 2*(sma - j);
	        }
	        cout << ans << " ";
	    }
	    cout << "\n";
	}
	return 0;
}
Editorialist's code (C++)
#include "bits/stdc++.h"
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
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;
		vector<int> a(n), b(n);
		for (int &x : a) cin >> x;
		for (int &x : b) cin >> x;

		int sumA = accumulate(begin(a), end(a), 0);
		int sumB = accumulate(begin(b), end(b), 0);
		const int inf = 1e9 + 7;
		vector dp(n+1, vector(sumB + 1, inf));
		vector<int> ans(n, inf);
		dp[0][0] = 0;
		for (int i = 0; i < n; ++i) {
			for (int ch = i+1; ch >= 1; --ch) {
				for (int j = b[i]; j <= sumB; ++j) {
					dp[ch][j] = min(dp[ch][j], dp[ch-1][j-b[i]] + max(0, a[i] - b[i]) - a[i]);
					if (j >= sumA) ans[n-ch] = min(ans[n-ch], 2*dp[ch][j]);
				}
			}
		}
		for (int i = 0; i < n; ++i) {
			if (ans[i] < 1e8) cout << ans[i] + 2*sumA << ' ';
			else cout << -1 << ' ';
		}
		cout << '\n';
	}
}
1 Like