Search code examples
pythonalgorithmdynamic-programmingproof-of-correctness

How to prove the correctness of this solution for Codeforces problem "A. Boredom"?

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.

Solution

  • Some observations:

    • The order of the sequence is of no importance. If you'd shuffle the input order, the output would still be the same. So we might as well think of the input being sorted.
    • When values are "isolated", i.e. the array no longer has a value that is one less or one more, then we should always take it: there is no "penalty" -- only the chosen value gets removed from the array. If we wouldn't take it, it would still be there after all other possible moves have been performed, and we would still take it at the end of the game then.
    • If the same value occurs multiple times, and we decide to take one of those, we should certainly also take the other occurrences, as then we arrive in the case of the previous point: there will be no more "collateral damage" (as the "neighboring" values were already removed by the first pick).
    • If we would have an input with 𝑛 numbers with all unique numbers between 1 and 𝑛, and 𝑛 would be odd, then we maximize by taking 1, 3, 5,... 𝑛-2, 𝑛.
    • If we would have the same as above, but now with 𝑛 even, we would take 2, 4, 6, ..., 𝑛-2, 𝑛. We could also have considered 1, 3, 5,... 𝑛-3, 𝑛-1, or leave an extra gap somewhere in between, (e.g. 1, 3, 6, 8, 10,... 𝑛-2, 𝑛), but in the end the first variant would have the maximum sum.
    • When moving through an input from left to right, we only have to consider two scenarios: the first one where the current value minus one was selected, and a second scenario where it was not selected. In the second scenario, we can select the current value, in the first scenario we can't, because it was removed by the previous move. If we select the current value using the second scenario, it automatically becomes the first scenario for the next potential move on (value plus one), while the first scenario becomes the second with regards to (value plus one). So we alternative scenarios and at the end pick the best scenario of both.

    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.

    Example visualisation

    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