How to aggregate the elements in a nested list

I have a list of lists as follows:

original_list = [['B_S', 'O', 'O', 'O'],
                 ['B_S', 'O', 'O', 'O'],
                 ['O', 'O', 'B_S', 'O'],

                 ['O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'O'],

                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'B_S', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'O', 'O', 'O']]

I want to aggregate each three list to be one list based on the majority of elements. That is, if two positions have the same element, the new list will take the same element in the same position. The desired output should be like:

desired_output = [['B_S', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'O', 'O', 'O']]

I have tried with this complex function:

def collect_labels(lab_user):

  def aggregate_labels(a,b,c):
     aggregated_list = []
     # iterate over 3 lists
     for x,y,z in zip(a, b, c):
        if x != 'O':
          aggregated_list.append(x)
        elif y != 'O':
          aggregated_list.append(y)
        #elif z != 'O':
          #aggregated_list.append(z)  # you can improve the code
        else:
          aggregated_list.append(z)
     return aggregated_list

  result = [aggregate_labels(lab_user[i], lab_user[i+1], lab_user[i+2]) for i in range(0,len(lab_user)-2, 3)] 
  return result

But it return the wrong result:

wrong_result= [['B_S', 'O', 'O', 'O'],
               ['O', 'O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'O'],
               ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'B_S', 'O', 'O']]

I will be very grateful if you could help me to correct my code!

Answer

A broader solution using loop:

def merge_result(data):
    ar = []
    for i in range(0, len(data)-2, 3):
        temp = []        
        for j in range(len(data[i])):
            if data[i][j] == data[i+1][j]:
                temp.append(data[i][j])
            elif data[i+2][j] == data[i+1][j]:
                temp.append(data[i+1][j])
            elif data[i][j] == data[i+2][j]:
                temp.append(data[i][j])
            else:
                temp.append('O')
        ar.append(temp)
    return ar

if __name__ == "__main__":
    original_list = [['B_S', 'O', 'O', 'O'],
                 ['B_S', 'O', 'O', 'O'],
                 ['O', 'O', 'B_S', 'O'],

                 ['O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'O'],

                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'B_S', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'O', 'O', 'O'],
                 ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'O', 'O', 'O']]
    print(merge_result(original_list))

Output:

[['B_S', 'O', 'O', 'O'], ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O'], ['O', 'O', 'O', 'O', 'O', 'B_S', 'O', 'O', 'B_S', 'O', 'O', 'O']]

Leave a Reply

Your email address will not be published. Required fields are marked *