I'm trying to build a simple contact tracing system in Python. The following function should get 2 persons in my database and return if they've had (in)direct contact, and if indirect it returns the amount of people in between them.

contact_tracing_datastructure = [{'Koen', 'Pieter'}, {'Kim', 'Koen', 'Bart'}, {'Tom', 'Pieter'}, {'Yana', 'Bart'}, {'Wouter', 'Yana'}, {'Wouter', 'Bert'}, {'Koen'}]

already_checked = set()

def x_had_contact_with_y_with_distance(meetings, person_x, person_y, distance=0):
    contacts_x = give_contacts(meetings, person_x)
    already_checked.add(person_x)
    if person_y in contacts_x:
        return True, distance
    elif len(already_checked) != len(population):
        distance  = 1
        for i in contacts_x:
            found = x_had_contact_with_y_with_distance(meetings, i, person_y, distance)
            if found:
                return True, distance
    else:
        return False

print(x_had_contact_with_y_with_distance(contact_tracing_datastructure, "Koen", "Wouter")

The function keeps going however until maximum recursion depht is reached, even though I've added the already_checked variable to make it so it doesn't check a person twice. The function should return (True, 2) in this case. Give_contacts(person_x) simply gives a set of people person_x has been in contact with.
How can i stop the recursion when a solution has been found? Thanks in advance!

CodePudding user response：

reshape

First I would reshape your data from list(set(str)) to dict(str,set(str)) -

# list of set of string
input = [{'Koen', 'Pieter'}, {'Kim', 'Koen', 'Bart'}, {'Tom', 'Pieter'}, {'Yana', 'Bart'}, {'Wouter', 'Yana'}, {'Wouter', 'Bert'}, {'Koen'}]

# dict of (string -> set of string)
index = dict()
for s in input:
  for name in s:
    if name in index:
      index[name] = index[name] | s - {name}
    else:
      index[name] = s - {name}

for (name, others) in index.items():
  print(name, others)

Pieter {'Koen', 'Tom'}
Koen {'Pieter', 'Bart', 'Kim'}
Kim {'Koen', 'Bart'}
Bart {'Yana', 'Koen', 'Kim'}
Tom {'Pieter'}
Yana {'Bart', 'Wouter'}
Wouter {'Yana', 'Bert'}
Bert {'Wouter'}

graph

The superior index shape allows you to easily write graph(t, q) -

def graph(t, q):
  def loop(r, q):
    yield r
    for name in t[q]:
      if name not in r:
        yield from loop((*r, name), name)
  return loop((q,), q)

for line in graph(index, "Koen"):
  print(line)

('Koen',)
('Koen', 'Bart')
('Koen', 'Bart', 'Kim')
('Koen', 'Bart', 'Yana')
('Koen', 'Bart', 'Yana', 'Wouter')
('Koen', 'Bart', 'Yana', 'Wouter', 'Bert')
('Koen', 'Pieter')
('Koen', 'Pieter', 'Tom')
('Koen', 'Kim')
('Koen', 'Kim', 'Bart')
('Koen', 'Kim', 'Bart', 'Yana')
('Koen', 'Kim', 'Bart', 'Yana', 'Wouter')
('Koen', 'Kim', 'Bart', 'Yana', 'Wouter', 'Bert')

trace

Utilizing graph we can easily write trace(t, q, p) -

def trace(t, p, q):
  for line in graph(t, p):
    if line[-1] == q:
      yield line

Let's see the ways Koen traces to Kim -

for contact in trace(index, "Koen", "Kim"):
  print(len(contact) - 1, "->".join(contact))

2 Koen->Bart->Kim
1 Koen->Kim

And how Bert traces to Kim -

for contact in trace(index, "Bert", "Kim"):
  print(len(contact) - 1, "->".join(contact))

5 Bert->Wouter->Yana->Bart->Koen->Kim
4 Bert->Wouter->Yana->Bart->Kim

And how Yana traces Wouter -

for contact in trace(index, "Yana", "Wouter"):
  print(len(contact) - 1, "->".join(contact))

1 Yana->Wouter

resilience

Modify graph so that any names can be searched, even names that don't appear in the graph -

def graph(t, q):
  def loop(r, q):
    try:                                       # <- try
      yield r
      for name in t[q]:
        if name not in r:
          yield from loop((*r, name), name)
    except KeyError:                           # <- name not found
      return                                   # <- stop iteration
  return loop((q,), q)

for contact in trace(index, "Vivian", "Sylvia"):
  print(len(contact) - 1, "->".join(contact))

(no output)