Contact Tracing Function

Question

I'm doing a past-year lab test in preparation for my upcoming test and I came across this question that took me a while to solve.

It asks the programmer to create a function, trace_contacts(), that takes in 2 parameters (patient, history), and outputs a list of people who may have been infected (either directly or indirectly from the initial patient).

Parameters:
patient - Contains a name in string format
history - Contains a list of recent interactions

For example:

patient = "Jason"    

history = [
("Jason", "Gideon", -3),
("Zac", "Yacob", -3),
("Gideon", "Brian", -1),
("Cindy", "Gideon", -2),
("Darren", "Jason", -5),
("Jason", "Vivian", -6)
]

The history list is a list of tuples. Each tuple has 3 sets of values, the first 2 strings contain the names of the 2 people who met. The 3rd index repesents the day they met (in relation to the day the initial patient was diagnosed, which is day 0). See below for a hopefully clearer visualization haha.

catches the virus	not yet infectious	infectious period	develops symptoms & stops meeting people
-7 (7 days before Day 0)	-6	-5 to -1	Day 0

The trace_contacts() function should first check the history for any interactions with the initial patient during his infectious period. It should ignore any meeting before the infectious period.

It should then check, of those who met the patient during his infectious period, if they had met anyone else during their infectious period (2-6 days after their first interection with the initial patient). Again, interactions 0-1 days after them meeting the initial infected should be ignored (because they are not yet infectious).

The function should then keep looping, to check for the next layer of indirect cases. And only stop and output the list, of people at risk, when there are no more indirect interactions.

Notes:

1. There should be no duplicate names in the output list.
2. The names in each tuple within history can be in any order. The function should be able to find risky interactions no matter where the infected patient is in the tuple.
3. The code should be able to be written within 15-20mins (average time/qn in my lab test haha)

The code I've written is as follows: Just wondering if anyone has any, more efficient way to write code for this question? Thanks in advance!

def trace_contacts(patient, history):
    caught_virus_list = []

    for meet in history:
        per_details = []
        if (patient in meet) and (meet[2]>-6):
            if meet[0] == patient:
                per_details.append(meet[1])
            elif meet[1] == patient:
                per_details.append(meet[0])

            per_details.append(meet[2])

            caught_virus_list.append(per_details)

    store_list = caught_virus_list[:]
    first_round = True

    while store_list != caught_virus_list or first_round:
        first_round = False
        store_list = caught_virus_list[:]

        #sort earliest meet date
        caught_virus_list.sort()

        prev_pax = ""
        sorted_caught_virus_list = []

        for pax in caught_virus_list:
            if pax[0] != prev_pax:
                sorted_caught_virus_list.append(pax)
            else:
                continue

            prev_pax = pax[0]

        caught_virus_list = sorted_caught_virus_list

        infect_names = [pax[0] for pax in caught_virus_list]
        infect_date = [(pax[1]+2) for pax in caught_virus_list]
        infect_names.insert(0,patient)
        infect_date.insert(0,-5)

        for meet in history:
            output_details = []
            name1 = meet[0]
            name2 = meet[1]
            date = meet[2]
            
            if (name1 in infect_names) and (name2 not in infect_names):
                date_idx = infect_names.index(name1)
                if date >= infect_date[date_idx]:
                    output_details.append(name2)
                    output_details.append(date)
                    caught_virus_list.append(output_details)
            elif (name2 in infect_names) and (name1 not in infect_names):
                date_idx = infect_names.index(name2)
                if date >= infect_date[date_idx]:
                    output_details.append(name1)
                    output_details.append(date)
                    caught_virus_list.append(output_details)
        
    infect_names.pop(0)

    return infect_names

Answer 1

I think you went at it from the wrong direction. While I am not completely sure how your algorithm works and if everything in it is necessary, I think you can generify it a bit.

What the problem seems to require is simulation of the days passed. You have a set of infectious people, a set of people that will become infectious and a list of people that have been in touch with infectious people.

What you then do is update the list of infectious people at the start of the day and add the people who got infected to the list from which you update.

def find_newly_infectious_people(day: int, infected: set):
    return filter(lambda person: person[1] == day, infected.copy())


def add_newly_infectious_people(new_infectious, infectious):
    for i in new_infectious:
        infectious.add(i[0])


def remove_newly_infectious_people(new_infectious, infected):
    for i in new_infectious:
        infected.remove(i)


def update_infectious(day, infectious, infected):
    new_infectious = find_newly_infectious_people(day, infected)
    add_newly_infectious_people(new_infectious, infectious)
    remove_newly_infectious_people(new_infectious, infected)


def find_contacts_on_day(day, history):
    return filter(lambda contact: contact[2] == day, history)


def check_for_contact_with_infected_person(person1, person2, patient, infectious, infected, result):
    if person1 in infectious and person2 != patient:
        infected.add((person2, day + 2))
        result.add(person2)


def trace_contacts(patient, history: list):
    infectious = set()
    infected = {(patient, -5)}
    result = set()

    for day in range(-6, 1):
        update_infectious(day, infectious, infected)

        contacts = find_contacts_on_day(day, history)
        for contact in contacts:
            check_for_contact_with_infected_person(contact[0], contact[1], patient, infectious, infected, result)
            check_for_contact_with_infected_person(contact[1], contact[0], patient, infectious, infected, result)

    result = list(result)
    result.sort()

    return result