枚举和映射的使用

Question

newResidues = [1, 2, 3, 4, 5]
newI = [[1,0,1,0,1],[1,1,0,0,0],[1,0,0,1,0]]
newI = np.array(newI)

for i in range(newI.shape[0]):
    indices = [for i,num in enumerate(newI) if num == 1] #line1
    indicesToResidues = dict(zip(indices,newResidues[indices]))#line2

So I know my code is incorrect... this is just an attempt of what I'm trying to do...

in line1 what I'm trying do is make a list of the indices wherever there is a 1 in the matrix for each row, separately. So indices for the first row of the matrix would look something like [0,2,5]

in line 2 trying to map these indices to a particular value with the same index in a list newResidues.

Any help/comments will be appreciated, thanks!

Answer 1

Replace your loop with the following:

from itertools import compress
for i in newI:
    print map(lambda _: newResidues.index(_), compress(newResidues, i))

itertools.compress is pretty neat and seems perfect for your use case. After the above gets the values back from compress it finds the index of each value from newResidues -- this won't work correctly if you have duplicate values in newResidues however!