Search code examples
pythonlistmatrixlist-comprehension

Do an AND operation between elements of lists in a list (element by element)


The problem I'm facing is that I have a list of lists.

a = [[1, 1, 1, 0, 0, 0],
    [1, 1, 1, 0, 1, 1], 
    [1, 0, 0, 1, 0, 0]]

I want to do AND/OR/Majority-voting operations on the elements of the lists inside the list a. I want to get a single list result by, for example, doing an AND operation index by index so it see if elements on index 0 (in the lists) fulfill AND operation criteria so it append 1 to a result list incase all elements on index 0 are 1, else 0.

and_operation_on_a = [1, 0, 0, 0, 0, 0]
def and_op(alist):
    
    result_list = []
        
    for x, y, z in zip(alist[0], alist[1], alist[2]):
        if x == 1 and y == 1 and z == 1:
            result_list.append(1)
        else:
            result_list.append(0)

    
    return result_list

I tried something like this as the list has 3 lists inside and it works for me but what I want to do is to generalize it. For example, for the function to do an AND/OR/Majority-voting operation (element-by-element or index-by-index) if given a variable number of lists inside. I shall be really thankful to you If someone can help me. I am stuck badly.


Solution

  • Here's one way to generalize your code. This script is written so that in the event of a tie, the result for maj_op defaults to zero.

    a = [[1, 1, 1, 0, 0, 0],
        [1, 1, 1, 0, 1, 1], 
        [1, 0, 0, 1, 0, 0]]
    
    def and_op(alist):
        return_list = []
        for tup in zip(*alist):
            return_list.append(int(all(tup)))
        return return_list
    
    def or_op(alist):
        return_list = []
        for tup in zip(*alist):
            return_list.append(int(any(tup)))
        return return_list
    
    def maj_op(alist):
        return_list = []
        half = len(alist)//2
        for tup in zip(*alist):
            return_list.append(int(sum(tup)>half))
        return return_list
    
    print(and_op(a))
    print(or_op(a))
    print(maj_op(a))
    

    The result:

    [1, 0, 0, 0, 0, 0]
    [1, 1, 1, 1, 1, 1]
    [1, 1, 1, 0, 0, 0]
    

    Alternatively, this can all be shortened with list comprehension. In particular, the rewritten functions

    def and_op(alist):
        return [int(all(tup)) for tup in zip(*alist)]
    
    def or_op(alist):
        return [int(any(tup)) for tup in zip(*alist)]
    
    def maj_op(alist):
        return [int(sum(tup)>len(alist)//2) for tup in zip(*alist)]
    

    yield exactly the same results.


    Regarding your comment:

    def maj_op_2(a,b):
        return [int(sum(x*w for x,w in zip(col,b))>=0) for col in zip(*a)]