horses.py

import random


class Horse:
    def __init__(self, speed) -> None:
        self.speed = speed


class RaceTrack:
    size = 5

    def race(cls, horses):
        if len(horses) > cls.size:
            raise ValueError(f"Cannot race more than {cls.size} horses")
        horses.sort(key=lambda h: h.speed, reverse=True)

    def cull(cls, horses: list[Horse]):
        while len(horses) > 3:
            horses.pop()


horses = {Horse(speed=random.random() * 100) for _ in range(25)}

actual_top_3 = sorted(horses, key=lambda h: h.speed, reverse=True)[:3]


sorted_batches = []
reference_horse = None
current = None
while horses:
    if not current:
        batch = [horses.pop() for _ in range(RaceTrack.size)]
    else:
        reference_horse = current[-1]
        batch = [reference_horse] + [horses.pop() for _ in range(RaceTrack.size - 1)]

    RaceTrack.race(batch)
    RaceTrack.cull(batch)

    # We know that there are two horses faster than this one,
    # so if any horse is slower than the reference, they're out.
    if reference_horse in batch:
        batch = batch[batch.index(reference_horse) + 1 :]

    if batch:
        sorted_batches.append(batch)

assert len(sorted_batches) <= 5

by_fastest_horse = {batch[0]: batch for batch in sorted_batches}

sorted_fastest = RaceTrack.race(by_fastest_horse.keys())
survivors = RaceTrack.cull(sorted_fastest)

by_fastest_horse = {
    batch[0]: batch for batch in sorted_batches if batch[0] in survivors
}