Traba: Given a list of workers and who they were referred by, find the worker that has referred the most workers

worker_referrals.py

# ### Context

# At Traba, one way that we've grown our worker pool is by offering incentives for our existing workers to refer people in their network. We'd like to do some analysis on which workers have been responsible for the most worker acquisition.

# ### Prompt

# Given an array of worker id pairs, where the 0th index represents a worker on Traba and the 1st index, if defined, represents the worker that referred the aforementioned worker, write a function that determines which worker is responsible for the most referrals on Traba, direct or indirect.

# ### Examples

# ```
# input – workers = [
#   ['A', undefined],
#   ['C', 'A'],
#   ['B', undefined],
# ]
# output - 'A'


# input – workers = [
#   ['A', undefined],
#   ['C', 'A'],
#   ['D', 'C'],
#   ['B', undefined],
#   ['E', 'B'],
# ]
# output - 'A'

# input – workers = [
#   ['A', undefined],
#   ['C', 'A'],
#   ['D', 'C'],
#   ['B', undefined],
#   ['E', 'B'],
#   ['F', 'B'],
# ]
# output - 'A' OR 'B'
# ```


class WorkerReferrals:
  def __init__(self):
    # Trackers the referral adjacency list as a Hashtable:
    # Referrer -> [People referred]
    self.referral_tracker = {}
    
    # Used to memoize number of total referrals per person
    self.referrers = {}
  
  def get_max_referrer_worker(self, worker_referrals):
    # Build a list of referrer to referrees
    for referral in worker_referrals:
      # If we find a new referrer, add it to list and add new entry as referree
      if referral[1] in self.referral_tracker:
        self.referral_tracker[referral[1]].append(referral[0])
      else:
        # Otherwise add new entry as referree
        self.referral_tracker[referral[1]] = [referral[0]]

    # Find the referrer with the most referrals
    cur_max = (None, 0)
    for referrer in self.referral_tracker:
      num_referrals = self.get_num_referrals(referrer)
      if cur_max[0] is None or num_referrals > cur_max[1]:
        cur_max = (referrer, num_referrals)

    return cur_max[0]
        
  def get_num_referrals(self, worker_id):
    if worker_id not in self.referral_tracker:
      return 0

    if worker_id in self.referrers:
      return self.referrers[worker_id]
      
    # Direct referrals
    res = len(self.referral_tracker[worker_id])

    # Indirect referrals
    for worker in self.referral_tracker[worker_id]:
      res = res + self.get_num_referrals(worker)

    self.referrers[worker_id] = res
    
    return res

  def clear(self):
    self.referral_tracker = {}
    self.referrers = {}


wr = WorkerReferrals()
assert wr.get_max_referrer_worker([
  ['A', None],
  ['C', 'A'],
  ['D', 'C'],
  ['B', None],
  ['E', 'B'],
  ['F', 'B'],
]) == "A"

assert wr.get_max_referrer_worker([
  ['A', None],
  ['C', 'A'],
  ['B', None]
]) == "A"

wr.clear()
assert wr.get_max_referrer_worker([
  ['A', None],
  ['C', 'A'],
  ['D', 'C'],
  ['B', None],
  ['E', 'B']
]) == "A"

wr.clear()
assert wr.get_max_referrer_worker([
  ['A', None],
  ['C', 'A'],
  ['B', None],
  ['E', 'B'],
  ['F', 'B']
]) == "B"

wr.clear()
assert wr.get_max_referrer_worker([]) is None

wr.clear()
assert wr.get_max_referrer_worker([
  ['A', None],
  ['B', None]
]) is None