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:

    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