I want to reduce time complexity of my code that solves the following problem.

The program receives a positive distance d followed by a sequence of n events: An event can be either that a person appeared at some position, or that a person disappeared from some position. After each event, your program should output "red" if there are currently at least two persons at distance exactly d, and "yellow" otherwise. Note that initially (before the first event) all positions are empty. The events are given as a 2D array event[0..n-1][0..1], where event[i][0] is "1" if a person appeared at position event[i][1] (on the x-axis), and "-1" if a person disappeared from that position. The output is an array light[0..n-1] of strings, where light[i] is the color directly after event[i].

I used the following code to solve this problem.

def cal_dis(d,pairs):
    found = False
    for i in range(0, len(pairs)-1):
        if(pairs[i+1]- pairs[i]) == d:
            found = True
    return "red" if found else "yellow"

def tracker(n,d,event):
    ret = [""] * n
    pairs = []
    for i in range(0,n):
        if event[i][0] == 1:
            ret[i]= cal_dis(d,sorted(pairs))
            ret[i]= cal_dis(d,sorted(pairs))
    return ret

Sample Input

first row: n d

next n rows: two integers; the first one is either -1 or 1, the second one is nonnegative (these n rows correspond to the array "event")

7 4
1 2
1 10
1 6
-1 2
-1 6
1 9
1 14

Sample Output

n rows each consisting of either the string "red" or "yellow".

  • 1
    \$\begingroup\$ In one place the input is called a two-dimensional array, and later it seems to be a file-like. Are the input and output stdin/stdout, or actually in-memory lists? \$\endgroup\$
    – Reinderien
    Aug 28, 2021 at 9:20

3 Answers 3


For a faster algorithm, update a count of neighboring pairs at distance D as you go.

from collections import defaultdict

def num_at_distance(d, pairs):
    # Keep track of whether each position is occupied or not.
    # Accepts a distance and a number of positions to update. Each position is a 2-tuple of:
    #   - -1 for not-occupied or 1 for occupied
    #   - A position which changed.
    # Returns [yellow, yellow, red...] as a generator

    # 'occupied' maps a position to whether it's occupied.
    # Everything is unoccupied to start (False).
    # For slightly faster performance you could switch to an array.
    occupied = defaultdict(bool)
    num_at_distance_d = 0
    for now_occupied, position in pairs:
        # A is the position D to the left of the update
        # i is the updated position
        # B is the position D to the right of the update
        a, i, b = position-d, position, position+d
        now_occupied = bool(now_occupied==1)

        if occupied[i] != now_occupied:
            occupied[i] = now_occupied
            if occupied[i]:
                # Use the fact that python's True and False are 1 and 0 for some math
                num_at_distance_d += occupied[a] + occupied[b]
                num_at_distance_d -= occupied[a] + occupied[b]
        if num_at_distance_d == 0:
            yield "yellow"
            yield "red"

if __name__=='__main__':
    input = 4, [(1,2), (1,10), (1,6), (-1, 2), (-1, 6), (1, 9), (1,4)]
    output = list(num_at_distance(*input))
  • \$\begingroup\$ Works perfectly! \$\endgroup\$ Aug 29, 2021 at 23:15

Use understandable names

cal_dis? Do you mean calculate_distance? But it doesn't calculate. Maybe find_at_distance or just find_distance will be better? And d is also too short (but if all the project is these two functions, it may ok). Maybe dist could be better? And event is a list, not a single event. Should it be called events? And pairs are positions, right?

Some simplifications

In cal_dis - when you find the element at distance d, you don't need to break the loop - you just need to return "red", and return "yellow" after the loop. Drop the found flag.

More Pythonic way to write this is something like

return "red" if any(b - a == d for a, b in zip(pairs, pairs[1:])) else "yellow"

but I'm not sure it's more readable, especially for beginners.

In tracker - you don't need to create ret full of empty values, you can append elements unto it. The last line of both if-else branches is the same, so you can move if out of if-else:

if event[i][0] == 1:
ret.append(d, sorted(pairs))

Now, do we need i? It is used only to address elements of event list, so we can iterate this list:

for action, position in events:
    if action==1:

All together:

def find_distance(dist, positions):
    positions = sorted(positions)
    for i in range(len(positions) - 1):
        if positions[i+1] - positions[i] == d:
            return "red"
    return "yellow"

def tracker(_, dist, events): #first argument is excessive, we have len(events)
    result = []
    positions = []
    for action, position in events:
        if action == 1:
        result.append(find_distance(dist, positions))
    return result

Better algorithm ideas

  • you can track current state. If there was a removal, the result can't switch from "yellow" to "red", and if data was added, "red" can't switch to "yellow".

  • if all positions are unique, you may keep them in a set instead of list, if they are not - use collections.Counter. This way you'll avoid sorting the positions every event, and checking if addition switched to "red" will be done in one check:

    if (position - dist) in positions or (position + dist) in positions:

  • to efficiently check if the removal switches state to "yellow" you should keep track on every existing pair with distance dist in the set of pairs. Doing this for non-unique positions (with Counter) will be a bit tricky.


IMO this is an algorithms practice question, and I don't think you want a review of your code. You just want the answer. Since this is doubtless from a practice site, why not just read the sample answers there?

That said, here is a hint on the answer: COUNT the number of pairs of people at distance exactly d as you go.

Whenever you update whether a person is at a particular location, check to the left d places, and to the right d places. If it helps, you can pad the array with d empty places to the left and right when you start.

This reduces the running time to linear, instead of quadratic.

  • \$\begingroup\$ can you point me to the practice website that you're talking about? \$\endgroup\$ Aug 28, 2021 at 11:24
  • \$\begingroup\$ Where did you get this problem? That's the website I mean. \$\endgroup\$ Aug 29, 2021 at 0:04
  • \$\begingroup\$ I didn't get it on a website for God's sake. A friend sent it to me to help him solve it. I was able write the code in a few minutes but I ran out of ideas to optimize it. If you know a website that has this question, then I can just see the sample solution and that's it. Why would someone post it here if they already could see the solution on a website? \$\endgroup\$ Aug 29, 2021 at 0:34
  • \$\begingroup\$ Code review is /filled/ with competitive coders "why is my code too slow for this competitive algorithms question from a site". That, together with the highly unnatural nature of this problem, made me think you got it from a site. Now I think your friend got it from a site or as an interview question. If this is a just a problem one of you ran into, great job on writing it up in a very clear way! \$\endgroup\$ Aug 29, 2021 at 20:34
  • \$\begingroup\$ Thanks for understanding \$\endgroup\$ Aug 29, 2021 at 23:15

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.