create sublist of indices with each sublist referring to set of unique tuples from a list of tuples

Viewed 75

I am trying to create sub list of indices by grouping indices of tuples with any of the elements being common from a list of tuples or keeping unique tuples indices separate. The definition of unique tuple being no element of the tuple is same as the elements in same position of other tuples in the list. Example: List which groups same company together,with same company defined as same name or same registration number or same name of CEO.

company_list = [("companyA",0002,"ceoX"),
                ("companyB"),0002,"ceoY"),
                ("companyC",0003,"ceoX"),
                ("companyD",004,"ceoZ")]

The desired output would be:

[[0,1,2],[3]]

Does anyone know of a solution for this problem?

3 Answers

The companies form a graph. You want to create clusters from connected companies.

Try this:

company_list = [
  ("companyA",2,"ceoX"),
  ("companyB",2,"ceoY"),
  ("companyC",3,"ceoX"),
  ("companyD",4,"ceoZ")
]

# Prepare indexes
by_name = {}
by_number = {}
by_ceo = {}
for i, t in enumerate(company_list):
  if t[0] not in by_name:
    by_name[t[0]] = []
  by_name[t[0]].append(i)
  if t[1] not in by_number:
    by_number[t[1]] = []
  by_number[t[1]].append(i)
  if t[2] not in by_ceo:
    by_ceo[t[2]] = []
  by_ceo[t[2]].append(i)

# BFS to propagate group to connected companies
groups = list(range(len(company_list)))
for i in range(len(company_list)):
  g = groups[i]
  queue = [g]
  while queue:
    x = queue.pop(0)
    groups[x] = g
    t = company_list[x]
    for y in by_name[t[0]]:
      if g < groups[y]:
        queue.append(y)
    for y in by_number[t[1]]:
      if g < groups[y]:
        queue.append(y)
    for y in by_ceo[t[2]]:
      if g < groups[y]:
        queue.append(y)

# Assemble result
result = []
current = None
last = None
for i, g in enumerate(groups):
  if g != last:
    if current:
      result.append(current)
    current = []
    last = g
  current.append(i)
if current:
  result.append(current)
print(result)

Fafl's answer is definitely more performant. If you're not worried about performance, here is a brute-force solution that might be easier to read. Tried to make it clear with some comments.

def find_index(res, target_index):
    for index, sublist in enumerate(res):
        if target_index in sublist:
            # yes, it's present
            return index

    return None  # not present
        
def main():
    company_list = [
        ('companyA', '0002', 'CEOX'),
        ('companyB', '0002', 'CEOY'),
        ('companyC', '0003', 'CEOX'),
        ('companyD', '0004', 'CEOZ'),
        ('companyE', '0004', 'CEOM'),
    ]

    res = []

    for index, company_detail in enumerate(company_list):
        # check if this `index` is already present in a sublist in `res`
        # if the `index` is already present in a sublist in `res`, then we need to add to that sublist
        # otherwise we will start a new sublist in `res`
        index_to_add_to = None

        if find_index(res, index) is None:
            # does not exist
            res.append([index])
            index_to_add_to = len(res) - 1
        else:
            # exists
            index_to_add_to = find_index(res, index)
        
        for c_index, c_company_detail in enumerate(company_list):
            # inner loop to compare company details with the other loop
            if c_index == index:
                # same, ignore
                continue
            if company_detail[0] == c_company_detail[0] or company_detail[1] == c_company_detail[1] or company_detail[2] == c_company_detail[2]:
                # something matches, so append
                res[index_to_add_to].append(c_index)
                res[index_to_add_to] = list(set(res[index_to_add_to]))  # make it unique

    print(res)

if __name__ == '__main__':
    main()

Check this out, I tried a lot for it. May be I am missing some test cases. Performance wise I think its good. I have used set() and pop those which lie in one group.

company_list = [
  ("companyA",2,"ceoX"),
  ("companyB",2,"ceoY"),
  ("companyC",3,"ceoX"),
  ("companyD",4,"ceoZ"),
  ("companyD",3,"ceoW")
]
index = {val: key for key, val in enumerate(company_list)}
res = []
while len(company_list):
      new_idx  = 0 
      temp = []
      val = company_list.pop(new_idx)
      temp.append(index[val])
      while new_idx < len(company_list) :
            if len(set(val + company_list[new_idx])) < 6:
                  temp.append(index[company_list.pop(new_idx)])
            else:
              new_idx += 1
      
      res.append(temp)
            
print(res)
Related