Encode event results as colors

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
break
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:
pairs.append(event[i][1])
ret[i]= cal_dis(d,sorted(pairs))
else:
pairs.remove(event[i][1])
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".

yellow
yellow
red
red
yellow
yellow
red

• 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? Aug 28, 2021 at 9:20

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]
else:
num_at_distance_d -= occupied[a] + occupied[b]
if num_at_distance_d == 0:
yield "yellow"
else:
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))
print(output)

• Works perfectly! 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:
pairs.append(...)
else:
pairs.remove(...)
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:
pairs.append(position)
...


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:
positions.append(position)
else:
positions.remove(position)
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.