stop matching

jfk.py

"""
python3 jfk.py
"""
from collections import namedtuple

Stop = namedtuple('Stop', ['sid', 'seq'])

stops_data = [
    Stop(sid=0, seq=0),
    Stop(sid=1, seq=1),
    Stop(sid=2, seq=2),
    Stop(sid=3, seq=3),
    Stop(sid=4, seq=4),
    Stop(sid=5, seq=5),
    Stop(sid=6, seq=6),
    Stop(sid=7, seq=7),
    Stop(sid=1, seq=8),
    Stop(sid=0, seq=9),
]


def get_stops(route_stops):
    curr_idx = 0
    result = []
    for stop_id in route_stops:
        curr_idx, stop = next(((i, s) for i, s in enumerate(stops_data[curr_idx:], curr_idx) if s.sid == stop_id), (0, None))
        if stop is None:
            return []
        result.append(stop)

    return result


# No match
assert [s.seq for s in get_stops([100, 101])] == []
print('passed 1')

# Start at origin
assert [s.seq for s in get_stops([0, 1, 2])] == [0, 1, 2]
print('passed 2')

# Start in middle
result = [s.seq for s in get_stops([4, 5, 6, 7, 1, 0])]
assert result == [4, 5, 6, 7, 8, 9], f'result is {result}'
print('passed 3')

# Start at end, when id starts to repeat
result = [s.seq for s in get_stops([1, 0])]
assert result == [8, 9], f'result is {result}'
print('passed 4')

here is a definition for get_stops that passes all test cases, it uses the property that we must find an exact match on the whole sequence and not just fuzzy matches

def get_stops(route_stops):
    curr_start = 0
    start_stop_id = route_stops[0]
    route_length = len(route_stops)

    while curr_start < len(stops_data) + route_length:
        
        curr_start, start_stop = next(
            ((i, s) for i, s in enumerate(stops_data[curr_start:], curr_start) if s.sid == start_stop_id), 
            (0, None)
        )

        # no start stop is found
        if start_stop is None:
            return []

        # the start stop found doesn't have enough stops after for this sequence
        if len(stops_data) < curr_start+route_length:
            return []

        test_sequence = stops_data[curr_start:curr_start+route_length]
        if [ts.sid for ts in test_sequence] == route_stops:
            return test_sequence

        curr_start += 1