Problem originally from: https://train.nzoi.org.nz/problems/1311

Every light-bulb is initially off, and O(N^2) solutions are too slow.

Emma is a massive nuisance at school by always meddling around with the light switches. Each day, she goes to toggle a selection of the school's N light switches, numbered from 1 to N. If the selected switch is originally on, she turns it off, and if it is originally off, she turns it on.

Emma chooses her light switches in a very particular way. For each day i, she starts at the ki-th light switch and moves towards higher-numbered switches, toggling every pi-th switch until she runs out of switches.

After D days, Emma is caught and is expelled from her school. You need to find the final states of all of the school's light switches.

All switches are initially off, and you can assume that no one else touches the switches.

I'm just wondering what would be a good way to go about this problem? It's a harder variation of the classic light bulb problem which can be seen here: https://math.stackexchange.com/questions/11223/the-final-state-of-1000-light-bulbs-switched-on-off-by-1000-people-passing-by


The original problem being, you have 100 light bulbs, and you flip every 2nd light bulb, go back to the start, and flip every 3rth light bulb, ... all the way till you start doing every 100th. But instead, in this new problem, instead of doing 100 runs, you might only do something like 8 runs. And instead of doing every 5th light bulb starting from 0, which would be a simple 5*N linear equation.

You might do every 5th light-bulb starting from the 2nd.

I'm not really sure how to solve this problem but, a time-complexity of O(N^2) is too slow. i.e., you can't just simulate doing or have a for-loop for every run given.


The first line includes two space-separated integers: N, the number of light switches, and D, the number of days before Emma is expelled.

The following D lines each contain two space-separated integers representing ki and pi for each day.

D <= N <= 50000


In the sample case, on day 1, Emma turns on the 1st switch, and then every 2nd switch after that until she runs out of switches:


Then, on day 2, Emma starts by turning on the 2nd switch and continues by turning off the 5th switch:


Sample Input 1

5 2
1 2
2 3
  • 1
    $\begingroup$ Note, that $p_i$ are only up to 50. So you will be ok with an $O(N logN + D log D + N * (max(p_i))^2)$ solution. $\endgroup$ Jul 19, 2023 at 10:58
  • $\begingroup$ For simplicity, in computational complexity, this website tends to just assume N = D for the worst cases. How do I get that solution though, I've been trying at this for hours and hours 😭 I remember when this problem came out a couple of years ago as well, I couldn't get it back then either. $\endgroup$ Jul 19, 2023 at 11:15
  • 1
    $\begingroup$ Note, for $i$-th bulb you are interested in a number of days such that $i \geq k_j$ and $i - k_j = 0 (mod p_j)$. You can sort all days by $k_j$. Then iterate through bulb numbers in increasing order and maintain a 2D table of size $50^2$ such that in cell (a, b) there is a number of $j$-th such that $k_j$ is not greater than the index $i$ of the current bulb, $p_j = a$ and $i - k_j = b (mod a)$. Obviously you can update this table on each iteration in O(max(p_j)^2) time. And this table if sufficient to calculate aforementioned value for each bulb. $\endgroup$ Jul 19, 2023 at 11:52
  • $\begingroup$ Thanks for this, my solution is quite similar, except the table is size N * max(P). $\endgroup$ Jul 20, 2023 at 11:39

1 Answer 1


Thanks to Vladislav Bezhentsev, and another article I read online, I now have a solution which is O(N * max(p_i)).

The way this works, is, there might be 50k different combinations of K, and P given. But P is at most 50. So, previously we were doing 50k for loops simulating turning on/off light bulbs.

But now, we just have an array saved that gives us all possible K, and P combinations. Lets call it Query_array, with all values 0 initially.

Query_array is a 2D list which has dimensions [0..50] by [0 .. N] (both inclusive)

You also have a light_bulb array that's size N, keeping track which light-bulb is on/off (which we loop over at the end to print ON OFF ON or whatever). And all values are 0 initially.

if K, and P come up, save Query_array[P][K] += 1

You loop through all P values, and then through all K mod P values.

For some specific P, and K mod P. You start off with amount_of_equations_so_far = 0

And then you loop through all values between K mod P, and N,

Going... K, K + P, K + 2P, K + 3P... <= N

You go:

amount_of_equations_so_far += Query_array[P][K]

light_bulb[K] += amount_of_equations_so_far

amount_of_equations_so_far += Query_array[P][K + P]

light_bulb[K + P] += amount_of_equations_so_far

amount_of_equations_so_far += Query_array[P][K + 2P]

light_bulb[K + 2P] += amount_of_equations_so_far


All the way till you pass N.

So this is basically just, 1 for loop, for all values of K which have the same value mod P. Rather than, doing an individual loop for K, and then another for K + P. And recalculating over-and-over again.

My solution below, but it could be more efficient.

# https://train.nzoi.org.nz/problems/1311

N, D = [int(x) for x in input().split()]

switches = [0 for _ in range(N)]

something_started = []

for p in range(51):
    for k in range(N):
        # Shows nothing started at value k, using p step_length

for i in range(D):
    k, p = [int(x) for x in input().split()]
    # with jump value p, something started at k
    something_started[p][k-1] += 1

for p in range(1, 51):
    for k in range(p):
        how_many_are_mod_k = 0

        # Checks all values starting at value k, ending at N, checking all values that are mod K
        # Just like a smartly done for loop

        # Why is this faster?
        # You're not doing a for loop for every D value
        # There can be 50k p, and k combos given.
        # Instead of looping for each of them, if a lot of them are the same P value, all K mod P values are grouped together
        for light_bulb_i in range(k, N, p):
            # print(p, light_bulb_i)
            how_many_are_mod_k += something_started[p][light_bulb_i]
            switches[light_bulb_i] += how_many_are_mod_k

for light_bulb in switches:
    if light_bulb % 2 == 0:

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.