I am looking at the Codeforces problem A. Boredom:
Given a sequence 𝑎 consisting of 𝑛 integers. The player can make several steps. In a single step he can choose an element of the sequence (let's denote it 𝑎𝑘) and delete it, at that all elements equal to 𝑎𝑘+1 and 𝑎𝑘−1 must also be deleted from the sequence. That step brings 𝑎𝑘 points to the player.
I found the following code submitted by user Leeisateam:
input()
z = [0] * 7**6
for i in map(int, input().split()):
z[i] += i
a = b = 0
for i in z:
a, b = max(a, i + b), a
print(a)
I understand what this code is doing up until the final loop starts. We're creating a list z
of the form
[0 X count_in_sequence(0), 1 X count_in_sequence(1), ..., n X count_in_sequence(n)].
After that b
is assigned the value of a
and a
uses that (previous) value in the next iteration. I tried induction, but I still can't understand why this code would always work.
Some observations:
Now let's look at the code.
The code projects the values to their corresponding index, so that they are automatically visited in ascending order.
Then the variables a
and b
are introduced. They represent two alternative solutions of picking values before the current value i
is taken (or not). The one represented by a
is one that might have taken the value i-1
, and so if that would be done, the value i
would not be available anymore (as deleted by the previous pick). By consequence, we don't consider the a
scenario for picking i
. On the other hand b
represents a variant where we surely did not pick i-1
, and so we are sure that after scenario b
we can also pick value i
. So now after having considered i
, we have two scenarios: a
(which wasn't extended) and b + i
(note also that i
has all the duplicates summed up as we know we want to take all or nothing of it).
In the next iteration of the loop, the role of a
and b
change, since now b
has possibly selected that previous i
, which now is i-1
, and that previous b
is now a
(it is like a swap).
And so the loop "alternates" through the z
array, keeping track of two possible scenarios.
Describing this in words makes it quite verbose, but it does help to step through the code with a piece of paper using some limited inputs.
Let's say we have input with the following (value, frequency) pairs:
(1,7), (2,9), (3,6), (4,12), (5,11), (6,10), (7,16), (8,10), (9,10), (10,9)
Then we could visualise the values for a
and b
as the loop makes iterations:
z
will be equal to:
[0, 7, 9, 6, 12, 11, 10, 16, 10, 10, 9, 0, 0, 0, .... 0]
k | z[k] | i | a | b |
---|---|---|---|---|
0 | 0 | 0 | 0 | 0 |
1 | 7 | 7 | 7 | 0 |
2 | 9 | 18 | 18 | 7 |
3 | 6 | 18 | 25 | 18 |
4 | 12 | 48 | 66 | 25 |
5 | 11 | 55 | 80 | 66 |
6 | 10 | 60 | 126 | 80 |
7 | 16 | 112 | 192 | 126 |
8 | 10 | 80 | 206 | 192 |
9 | 10 | 90 | 282 | 206 |
10 | 9 | 90 | 296 | 282 |
11 | 0 | 0 | 296 | 296 |
... | ... | ... | ... | ... |
117648 | 0 | 0 | 296 | 296 |
i
represents what extra value can be added when we select the current value -- which is that value multiplied by its frequency in the input. This value is only allowed to be added to b
as the scenario for a
has allowed for the selection of the immediate predecessor value. Still, adding this value i
to b
is a potential new value for a
, but only if it is better than what a
is. In the example we see this is b+i
not better than a
when i
is 0 (which is true in the largest part of the z
list... having 76 entries).
And at the end a
has the solution: 296