Using numpy to inspect one array and create other fast-CodePudding

I have a simple vector that can only contains 1 , -1 or 0. It represents the direction of a particle. For instance array([1, 0]) or array([-1, 1]) I want to add a diagonal to component if there is none.

my current function is:

def addcomp1(vec: array):
    arr = np.empty(2, dtype=int)
    for i, el in enumerate((i if i else  1 for i in vec)):
        arr[i] = el
    return arr

This makes any straight vector that goes North, to go Nortwest. Or any that goes West, go northwest. I think I can make the function faster. My hints. I am only using a total of 4 inputs and 3 outputs. I am only using 1s, -1s and 0s. I would lru_cache the function but np.arrays are not hashable. Any ideas? Just use a dictionary that connects (1, 0) into (1, 1)? I would have to make them a hashable...

CodePudding user response：

If I understood correctly use:

import numpy as np


def addcomp1(vec: np.array):
    return np.where(vec == 0, 1, vec)


print(addcomp1(np.array([1, 1])))
print(addcomp1(np.array([0, 1])))
print(addcomp1(np.array([-1, 0])))

Output

[1 1]
[1 1]
[-1  1]

CodePudding user response：

I don't know how faster this could be, and I also don't know if this is exactly the type of answer you are looking for, nonetheless, I mention this simple line:

arr = [(lambda x: 1 if not x else x)(s) for s in vec]