Yandex programming contest: Alarms

Question

I've tried to solve this challenge about 4 hours during the contest, but all my attempts exceeded the time limit. I tried to solve it with min-heap, but can't pass all the tests. How it can be solved within given time and space requirements?

Problem statement

Programmer Alexey likes to work at night and does not like to come to work late. In order to precisely wake up in the morning, Alexey every evening starts \$N\$ alarm clocks on his phone. Each alarm clock is arranged in such a way that it rings every \$X\$ minutes from the time at which it was turned on.

For example, if the alarm clock was started at the moment of time \$t_i\$, then it will ring at the moments of time \$t_i, t_i + X, t_i + 2 * X\$ and so on. Moreover, if some two alarms begin to ring at one time, only one of them is displayed.

It is known that before waking up, Alexey listens every morning to exactly \$K\$ alarm clocks, and then wakes up. Determine the point in time when Alex wakes up.

Input format

Input format The first line contains three integers. \$N, X\$ and \$K\$ \$(1 ≤ N ≤10^5, 1≤X, K≤10^9)\$ - the number of alarms, the frequency of calls and the number of alarms that need to be turned off in order for Alex to wake up. The second line contains N integers - the points in time at which the alarm clocks were entered.

Requirements

Time limit: 2 seconds
Memory limit: 256Mb

Examples

Example 1

Input

6 5 10
1 2 3 4 5 6

Output 10

Example 2

Input

5 7 12
5 22 17 13 8

Output 27

Notes

In the second example, there are 5 alarm clocks with a frequency of 7 calls. For example, the first alarm clock will ring at times 5, 12, 19, 26, 33, etc. If you look at all the alarms at the same time, they will ring at the following times: 5, 8, 12, 13, 15, 17, 19, 20, 22 (2nd and 5th alarm clocks at the same time), 24, 26, 27, 29,…. On the 12th call Alexey must wake up, What corresponds to the point in time 27.

My solution

Classified as "time limit exceeded"

import heapq


def main():
    n, x, k = map(int, input().split(" "))
    times = list(map(int, input().split(" ")))

    heapq.heapify(times)

    answers = []

    while k:
        v = heapq.heappop(times)
        heapq.heappush(times, v + x)

        if len(answers) == 0 or answers[len(answers)-1] != v:
            answers.append(v)
            k -= 1

    print(answers[k-1])


if __name__ == "__main__":
    main()

Answer 1

Not all counting problems require enumerating the items being counted. If you were asked to count how many sheep there are in total if there are \$n\$ trucks with \$k\$ sheep each, you could write something like:

total_sheep = 0
for truck in range(n):
    for sheep in range(k):
        total_sheep += 1

Or you could cut to the chase and compute it as:

total_sheep = n * k

The problem you are trying to solve is more subtle, but can be attacked in a similar fashion. Modular arithmetic will be your friend in doing this. Rather than looking at the setting times as single numbers, convert them to tuples of (quotient, remainder) after dividing by \$X\$. This would e.g. convert the list of times for the second example into:

 5 --> (0, 5)
22 --> (3, 1)
17 --> (2, 3)
13 --> (1, 6)
 8 --> (1, 1)

We can use this information to prune the list of alarms: if any two alarms have the same remainder, they will ring at the same time, so we only need to keep the smaller one. This would convert the above list, after also sorting it, into:

 5 --> (0, 5)
 8 --> (1, 1)
13 --> (1, 6)
17 --> (2, 3)

Those two numbers tell us in which of the \$X\$ minute periods the alarm sounds for the first time, and at what offset into that period does it first go off. So we can process it sequentially and know that:

• at the end of the first (0) period, 1 alarm has sounded for the first time, a total of 1 alarm will sound in every subsequent period, and the total number of alarms sounded is also 1.
• at the end of the second (1) period, 2 alarms have sounded for the first time, 3 alarms will sound in every subsequent period, and the total number of alarms sounded is 4.
• at the end of the third (2) period, 1 new alarm has sounded, 4 alarms will sound every period, and a total of 8 alarms will have sounded.

Since there are no more alarms to process, we have 8 alarms so far, 4 more sounding each period, and want to reach a total of 12, some simple math tells us that this will happen during the fourth (3) period, and that it will be the last of the alarms that will reach it.

To make the math more clear, lets imagine that we wanted to reach a total of 14 alarms instead. Since 8 have already sounded, we have 14 - 8 = 6 more to go. Since 4 alarms will sound in each period, and the quotient and remainder of 6 divided by 4 are 1 and 2, we know that we will reach our target after 1 full more period, plus 2 of the four alarms in the next period. This translates to 4 full periods, plus the time for the second alarm in that period to sound. The time of 4 full periods is 7 * 4 = 28. We need to add the offset of the second alarm when sorting them by offset, not by start time, so we need to add 3, not 1, and the end result would be 31.

This general outline omits many gruesome details, like what to do if we reach the target number of alarms before we finish processing the list. And the math described above is hard to get right in all corner cases, as there are many chances for off-by-one errors. I'm not going to spoil the fun for you by coding it up: give a try, and I'll be happy to review your next version.

Answer 2

The minheap approach seems correct. The problem is in the answers. It may grow quite large (up to \$10^9\$). All the reallocations due to its growth are very costly.

However you don't need it at all. You only care about the time of the most recent alarm. Just one value:

    while k:
        v = heapq.heappop(times)
        heapq.heappush(times, v + x)

        if last_alarm != v:
            last_alarm = v
            k -= 1

That said, an idiomatic way to access the last element of the list is answers[-1].