Search code examples
pythonlist-comprehensionnested-loops

How do I convert this nested for loop to a map or list comprehension in python?


Note: this is a working program, I just need to make it faster. This is the nested for loop:

taf = []
for i in range(len(piou)):
    counter = 0
    for j in range(len(piou[i])):
        
        if piou[i][j] == True:
            counter = counter + 1
    if counter == len(piou[i]):
        taf.append(i)

I tried doing this:

taf = [i for i in range(len(piou))
       if counter == len(piou[i])
       for j in range(len(piou[i]))
       if piou[i][j] == True]

but I can't figure out how to put the counter variable in between along with its increment.

Note: piou is a 2D tensor with boolean values and I am trying to find and append all the rows (i) index values to taf list which have all "true" column values.

piou looks like this:

tensor([[ True,  True,  True,  True, False, False],
        [ True,  True,  True,  True,  True,  True],
        [False, False,  True,  True,  True,  True],
        [ True,  True, False, False,  True,  True],
        [ True,  True,  True,  True,  True,  True]], device='cuda:0')

Solution

  • A loop where you're just incrementing a number can be replaced with sum(), but where that summation is just going to be compared with the length of the list, use all().

    As well, avoid for i in range(len(seq)) ... i, seq[i], use for i, x in enumerate(seq) instead.

    taf = [i for i, row in enumerate(piou) if all(row)]