Who can explain how to loop add an element to the numpy array by condition?

I wrote some code that should do add element 2 if i element of array A is 0 and add element 1 if i element of array A is not 0.

Here is the code itself:

import numpy as np
def finalconcat(somearray):
  for i in somearray:
    arraysome=[]
    if somearray[i]==0:
      arraysome=np.append(arraysome,[2],axis=0)
    else:
      arraysome=np.append(arraysome,[1],axis=0)
  return arraysome

Let me give you an example:

A=np.array([1,0,2,3,4,5])
C=finalconcat(B)
print(C)

It should come out:

[1,2,1,1,1,1]

But it comes out like:

[1.]

Please explain to me what is wrong here, I just don't understand what could be wrong...

CodePudding user response：

You have several issues:

1. arraysome=[] is inside your loop so for each iteration of somearray you are emptying arraysome. Consequently, you can never end up with more than one element in arraysome when you are all done.
2. You have for i in somearray. On each iteration i will be the next element of somearray; it will not be iterating indices of the array. Yet later you have if somearray[i]==0:. This should just be if i==0:.
3. If you want the resulting elements of arraysome to be integers rather than floats, then you should initialize it to be a an numpy array of integers.
4. You have C=finalconcat(B), but B is not defined.
5. You should really spend some time reading the PEP 8 – Style Guide for Python Code.

import numpy as np

def finalconcat(somearray):
    arraysome = np.array([], dtype=np.int)
    for i in somearray:
        if i == 0:
            arraysome = np.append(arraysome, [2], axis=0)
        else:
            arraysome = np.append(arraysome, [1], axis=0)
    return arraysome

a = np.array([1, 0, 2, 3, 4, 5])
c = finalconcat(a)
print(c)

Prints:

[1 2 1 1 1 1]