Find index in a case of 2D array-CodePudding

I want to find 2D array's index and make it array. for example:

data_pre=[[1,1,1,0,0,0],[1,0,1,0,0,0],[1,0,0,0,1,0],[1,0,0,0,0,0]]

i wanna find index that have one and wanna make it like this

b=[[0,1,2],[0,2],[0,4],[0]]

Code:

result = []
for i in range(len(data_pre)):
    arr=data_pre[i]
    currentArrResult=[]
    for j in range(len(arr)):
        if arr[j]==1:
            currentArrResult.append(j)
            result.append(currentArrResult)

I tried like that but output is wrong.

[[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 2], [0, 2], [0, 4], [0, 4], [0]]

I don't know which part is wrong...

CodePudding user response：

you should not collect output inside the inner loop. that may get a result like this :

[[0],[0, 1],[0, 1, 2],
[0],[0, 2],
[0],[0, 4],
[0]]

you can check that by printing currentArrResult after append finish.
but you got different outcome because of data reference.

you should collect result after inner loop finish its work.
like:

result = []
for i in range(len(data_pre)):
    arr=data_pre[i]
    currentArrResult=[]
    for j in range(len(arr)):
        if arr[j]==1:
            currentArrResult.append(j)
           #print(currentArrResult)
    result.append(currentArrResult)