from statistics import mean
from typing import List
import sys

DEFAULT_WEIGHT = 5


def read_data(file_name: str, std_in=True):
    satisfactions_list = list()
    file = sys.stdin if std_in else open(file_name, 'r')
    days = int(file.readline().rstrip())
    satisfaction_decreases = [int(c) for c in file.readline().rstrip().split(" ")]
    for i in range(days):
        satisfactions_list.append([int(c) for c in file.readline().rstrip().split(" ")])
    if not std_in:
        file.close()
    return days, satisfaction_decreases, satisfactions_list


class Contest:
    def __init__(self, _type: int, satisfactions: List[int], satisfaction_decrease: int, weight=DEFAULT_WEIGHT):
        self.times_on_bench = list()
        self.last_time_on_bench = 0
        self.weight = weight

        self.dissatisfaction = 0
        self.type = _type
        self.satisfactions = satisfactions
        self.satisfaction_decrease = satisfaction_decrease

    def __getitem__(self, key: int):
        return self.satisfactions[key - 1]

    def __str__(self):
        return str(self.type)

    def __repr__(self):
        return f"Contest(_type={self.type}, satisfactions={self.satisfactions}, satisfaction_decrease={self.satisfaction_decrease})"

    def heuristic(self, day):
        # return self[day] - self.weight*(self.dissatisfaction - self.satisfaction_decrease)
        return self[day] - self.weight*self.dissatisfaction

    def winner(self):
        self.dissatisfaction = 0

        self.times_on_bench.append(self.last_time_on_bench)
        self.last_time_on_bench = 0

    def loser(self):
        self.dissatisfaction -= self.satisfaction_decrease

        self.last_time_on_bench += 1

    def make_weight(self):
        return mean(self.times_on_bench) if len(self.times_on_bench) else DEFAULT_WEIGHT


def get_best(contests: List[Contest], day: int):
    # improve this by optimistically making everyone a loser and then going back to undo that for the winner
    best = contests[0]
    for contest in contests:
        if contest.heuristic(day) > best.heuristic(day):
            best = contest
    return best


def main(file_name="input.txt", weight=False):
    contests = list()
    picks = list()

    days, satisfaction_decreases, satisfactions_list = read_data(file_name)

    for (_type, (satisfactions, satisfaction_decrease)) in enumerate(zip(zip(*satisfactions_list), satisfaction_decreases), start=1):
        contests.append(Contest(_type, satisfactions, satisfaction_decrease))

    for day in range(days):
        best = get_best(contests, day + 1)
        for contest in contests:
            if contest.type == best.type:
                contest.winner()
                picks.append(contest)
            else:
                contest.loser()

    if weight:
        weights = [contest.make_weight() for contest in contests]
        contests = list()
        picks = list()

        for (_type, (satisfactions, satisfaction_decrease, weight)) in enumerate(zip(zip(*satisfactions_list), satisfaction_decreases, weights), start=1):
            contests.append(Contest(_type, satisfactions, satisfaction_decrease, weight))

        # print(weights)

        for day in range(days):
            best = get_best(contests, day + 1)
            for contest in contests:
                if contest.type == best.type:
                    contest.winner()
                    picks.append(contest)
                else:
                    contest.loser()

    return "\n".join([str(pick) for pick in picks])


if __name__ == "__main__":
    print(main("input2.txt"))


Measured memory_bytes: 11596