PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: pols_agyi_pols
Tester & Editorialist: iceknight1093
DIFFICULTY:
2850
PREREQUISITES:
Dynamic programming
PROBLEM:
N students are in a circular hostel with M rooms.
Initially, student i is in room A_i.
For each i and j, you also know the value H_{i, j}: the happiness of student i if he were to stay in room j.
You can perform the following operation:
- Choose two students i and j.
Move i clockwise by X rooms and j counterclockwise by X rooms.
Find the maximum possible total happiness of the students after performing some moves.
EXPLANATION:
Working with pairs of students is somewhat hard.
Observe that we can instead work with single students and fix how many moves each one makes: it doesn’t really matter who is paired with who as long as the total number of counterclockwise moves equals the total number of clockwise moves — the end result will be the same and that’s all that matters.
This gives us something to start with.
Let \text{dp}[i][j] denote the maximum possible happiness if we’ve considered students 1, 2, \ldots, i; and there are j ‘unpaired’ moves.
For convenience, if j \gt 0 we say that j anticlockwise moves must be made, while j \lt 0 means -j clockwise moves must be made.
Then,
- If student i makes k clockwise moves, he will end up at room (A_i + k\cdot X)\bmod M.
Let h be his happiness at ending up in this room.
Then the value we obtain here is \text{dp}[i-1][j-k] + h. - If student i makes k anticlockwise moves, he will end up at room (A_i - k\cdot X)\bmod M.
Again, if h is the happiness of this, the overall value is \text{dp}[i-1][j+k] + h.
By considering all possible values of k, \text{dp}[i][j] can be computed as the maximum among the above.
The final answer is \text{dp}[N][0]: the best value after all students have been considered, and no excess moves remain.
While this formulation is correct, it has a large flaw: it’s too slow.
In fact, the number of moves that a single student can make might have to be very large, so we have to constrain this somehow.
For that, we use the cyclicity property we have.
We can make the following observations:
- First, it’s enough to consider \text{dp}[i][j] values for -M \lt j \lt M.
This is because, if j \geq M, student i can just make M counterclockwise moves to reduce j by M while still remaining at his initial position. Repeat this till j \lt M. - Secondly, for similar reasons, it’s enough to consider k (the number of moves made) to also be from 0 to M-1; since if we make more, the count can be reduced by M without affecting the result.
This immediately gives us a more tractable solution.
There are \mathcal{O}(N\cdot M) states of \text{dp}[i][j], and for each one we perform M transitions for \mathcal{O}(NM^2) in total, which should already be enough to solve the problem.
This is enough to get AC if implemented well, though further optimizations can be made to improve runtime:
- Instead of -M \lt j \lt M, it’s enough to consider 0 \leq j \lt M.
This is because j clockwise moves can also be simulated using M-j anticlockwise moves. - One nice property of repeatedly adding X on cycles like this is that we in fact obtain \gcd(M, X) disjoint cycles, each of length \frac{M}{\gcd(M, X)} and rather equally spaced out.
Each element, when X is added/subtracted to it, only moves within its own cycle.
This allows us to bring the complexity down to \mathcal{O}\left(N\cdot \left(\frac{M}{\gcd(M, X)}\right)^2\right ), though the worst case complexity (when \gcd(M, X) = 1) remains the same.
TIME COMPLEXITY
\mathcal{O}(N M^2) per testcase.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define ll long long
int main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
ll kitne_cases_hain;
kitne_cases_hain=1;
//freopen("input.txt","r",stdin);freopen("output.txt","w",stdout);
cin>>kitne_cases_hain;
while(kitne_cases_hain--){
ll n,m,x,total;
cin>>n>>m>>x;
ll a[n];
for(int i=0;i<n;i++){
cin>>a[i];
}
ll h[n][m];
for(int i=0;i<n;i++){
for(int j=0;j<m;j++){
cin>>h[i][j];
}
}
//cin>>total;
ll jump=__gcd(m,x);
ll t=m/jump;
ll dp[t];
dp[0]=0;
for(int i=1;i<t;i++){
dp[i]=-1e18;
}
ll room;
ll dp1[t];
ll f;
for(int i=1;i<=n;i++){
for(int j=0;j<t;j++){
dp1[j]=-1e18;
}
for(int j=0;j<t;j++){
room=a[i-1]-1;
for(int c=0;c<t;c++){
if(j+c>=t){
f=(j+c)%t;
}else{
f=j+c;
}
dp1[f]=max(dp[j]+h[i-1][room],dp1[f]);
room+=jump;
if(room>=m){
room%=m;
}
}
}
for(int i=0;i<t;i++){
dp[i]=dp1[i];
}
}
cout<<dp[0]<<"\n";
}
return 0;
}
Editorialist's code (Python)
import math
for _ in range(int(input())):
n, m, x = map(int, input().split())
a = list(map(int, input().split()))
h = [list(map(int, input().split())) for i in range(n)]
length = m//math.gcd(m, x)
dp1 = [- 10**18]*length
dp1[0] = 0
for i in range(n):
pos = a[i]-1
dp2 = [- 10**18] * length
for j in range(length):
for k in range(length):
where = (pos + k*x) % m
dp2[(j+k)%length] = max(dp2[(j+k)%length], dp1[j] + h[i][where])
dp1, dp2 = dp2, dp1
print(dp1[0])