Python Group Values in a Nested List At Corresponding Indexes

Given a list nested with numpy lists, I want to go through each index in each corresponding list and keep track of the count depending on the element, which would be stored into a single list.

The minimum runnable code example below better showcases the problem.

import numpy as np


lst = [
    np.array([1,0,1]),
    np.array([1,1,1]),
    np.array([2,2,1])
]
print(lst)

#Problem, take count of each corresponding index within list
#If the element == 2, subtract one. If it's 1, add 1. 0 has no value
#Make the output just one list with the calculated count for each index

#Expected output:
#lst = [1, 0, 3]
# 1 in index 0 because there's two 1s in each lists first index, but a 2 subtracts one from the count
# 0 in index 1 since there's one 1 which brings the count to 1, but a 2 subtracts one from the count
# 3 in index 2 since there's three 1's, adding 3 to the count.

I have looked into from collections import defaultdict and from itertools import groupby, but I couldn’t understand how to use that to do anything else but sort a nested list.

Apologies if this issue is within the two sources I provided above.

Answer

If all the inner arrays have the same length you could transform it into a 2-dimensional array to leverage numpy vectorization:

# transform into two dimensional array
arr = np.array(lst)

# set -1 where values are equals to 2
t = np.where(arr == 2, -1, arr)

# sum across first axis (each corresponding positions on the inners lists)
res = t.sum(axis=0).tolist()

print(res)

Output

[1, 0, 3]