D - アルバイトの計画 / Part-Time Job Planning 解説 by admin
gpt-6-astra-mediumOverview
We want to choose at most \(K\) days to work a part-time job, with an interval of at least \(D+1\) days between working days, to maximize the total salary.
We use dynamic programming (DP) where the state represents “up to which day we have considered” and “at most how many days we can work”, handling special cases with simpler methods.
Analysis
Simply choosing days with the highest salary does not work
When \(D>0\), choosing a certain day prevents us from choosing days close to it.
For example, when \(N=3,\ K=2,\ D=1\), and the salaries are \([6,10,6]\), if we pick day \(2\), which has the highest salary, the total is \(10\) yen. On the other hand, picking day \(1\) and day \(3\) yields a total of \(12\) yen.
Therefore, a simple greedy approach that prioritizes days with higher salaries cannot find the correct answer. Also, checking all combinations of working days is too slow because there are up to \(2^N\) possibilities.
Branching on whether to work on the last day
We define \(F_k(i)\) as follows:
- The maximum salary obtained from day \(1\) through day \(i\), working for at most \(k\) days.
For day \(i\), we consider two choices:
- Do not work: the maximum value is \(F_k(i-1)\)
- Work: the latest day we could have worked previously is on or before day \(i-D-1\), so the maximum value is \(F_{k-1}(\max(0,i-D-1))+A_i\)
Thus, the transition is as follows:
\(F_k(i)=\max\left(F_k(i-1),\ F_{k-1}(\max(0,i-D-1))+A_i\right)\)
The base cases are \(F_0(i)=0,\ F_k(0)=0\).
With this approach, each state can be computed in constant time, leading to a time complexity of \(O(NK)\).
The number of workable days has an upper bound due to the rest condition
Let \(step=D+1\). To work \(m\) days, even if we pack the working days as closely as possible, the last working day is at day \(1+(m-1)step\).
Therefore, the maximum number of days we can possibly work is:
\(M=\left\lfloor\frac{N+D}{D+1}\right\rfloor\)
If \(K\ge M\), simply satisfying the rest condition automatically satisfies the limit on the number of days. In this case, we can eliminate the number of days \(k\) from the DP state.
Algorithm
In the implementation, we divide the problem into the following \(3\) cases:
1. Case \(D=0\)
Since there is no rest period, any day can be chosen freely.
Because all salaries are positive, we sort the salaries in descending order and output the sum of the top \(K\) values.
2. Case \(K\ge M\)
Since we do not need to consider the limit on the number of days, we can use 1D DP.
Let \(dp[i]\) be “the maximum salary obtained from day \(1\) through day \(i\)”. Then:
\(dp[i]=\max\left(dp[i-1],\ dp[\max(0,i-step)]+A_i\right)\)
We compute this for \(i=1,2,\ldots,N\) in order and output \(dp[N]\).
In the code, best represents \(dp[i-1]\) before the update and \(dp[i]\) after the update.
3. Otherwise
We compute \(F_k(i)\) as described in the analysis. However, it is not necessary to store the results for all \(k\).
To compute \(F_k(i)\), we only need:
- The value from the previous stage: \(F_{k-1}(i-step)\)
- The value obtained just prior in the current stage: \(F_k(i-1)\)
Therefore, we can reuse arrays to save memory.
First, compute the case of working at most \(1\) day
When working at most \(1\) day, we do not need to worry about the rest condition, so:
\(F_1(i)=\max(A_1,A_2,\ldots,A_i)\)
We store this in dp.
Update in the order of \(k=2,3,\ldots,K\)
Before updating, dp[i] represents \(F_{k-1}(i)\). The new results are stored in current.
Here, the earliest day on which it is possible to have worked \(k\) days is:
\(start=(k-1)step+1\)
For \(i<start\), it is impossible to work \(k\) days in the first place, so \(F_k(i)=F_{k-1}(i)\).
Therefore, we update as follows:
- Set
current = dp[:]to copy the previous stage. - Set
best = dp[start - 1]. - Iterate through \(i=start,\ldots,N\):
- Compute the value if we work on day \(i\):
dp[i - step] + A[i]. - Take the maximum of this value and
bestas the newbest. - Set
current[i] = best.
- Compute the value if we work on day \(i\):
- Set
dp = current.
Because best maintains \(F_k(i-1)\), this update directly corresponds to the transition formula described in the analysis.
Finally, output dp[N], which is \(F_K(N)\).
Complexity
- Time complexity:
- If \(D=0\): \(O(N\log N)\)
- If \(D>0\) and \(K\ge M\): \(O(N)\)
- Otherwise: \(O(NK)\)
- Space complexity: \(O(N)\)
Implementation Points
- The required interval is \(D+1\), not \(D\)
After working on day \(i\), the next day we can work is day \(i+D+1\). Setting step = D + 1 makes the index calculations easier to follow.
- DP tracks “at most \(k\) days”, not “exactly \(k\) days”
Depending on the rest condition, choosing a smaller number of days with high salaries might be more beneficial. We do not necessarily have to use up all days.
- Do not overwrite the array of the previous stage directly
dp[i - step] specifically requires the result for at most \(k-1\) days. Prepare current separately so that already-updated values do not get mixed in.
- A large integer type is required for the answer
The total salary can be up to \(5000\times10^9=5\times10^{12}\). While Python handles this automatically, a 64-bit integer should be used in other languages.
- In the branch for \(D=0\), the dummy element at the beginning is also sorted
Although A in the code includes a \(0\) for 1-based indexing, all \(A_i\) are positive and \(K\le N\), so this \(0\) will never be included in the top \(K\) elements when sorted in descending order.
Source Code
import sys
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
N, K, D = data[:3]
A = [0] + data[3:]
if D == 0:
A.sort(reverse=True)
print(sum(A[:K]))
return
step = D + 1
if K >= (N + D) // step:
dp = [0] * (N + 1)
best = 0
for i in range(1, N + 1):
value = A[i]
if i > step:
value += dp[i - step]
if value > best:
best = value
dp[i] = best
print(best)
return
dp = [0] * (N + 1)
best = 0
for i in range(1, N + 1):
if A[i] > best:
best = A[i]
dp[i] = best
for k in range(2, K + 1):
start = (k - 1) * step + 1
current = dp[:]
best = dp[start - 1]
for i in range(start, N + 1):
value = dp[i - step] + A[i]
if value > best:
best = value
current[i] = best
dp = current
print(dp[N])
if __name__ == "__main__":
main()
This editorial was generated by gpt-6-astra-medium.