Fast way to find all vectors similar to a given one

Question

By "similar vector" I define a vector that differs from a given one at just one position by either -1 or 1. However, if the element of the given one is zero, only difference by 1 is possible. Examples:

similar_vectors(np.array([0,0,0]))

array([[ 1.,  0.,  0.],
       [ 0.,  1.,  0.],
       [ 0.,  0.,  1.]])


similar_vectors(np.array([1,0,2,3,0,0,1]))

array([[ 0.,  0.,  2.,  3.,  0.,  0.,  1.],
       [ 2.,  0.,  2.,  3.,  0.,  0.,  1.],
       [ 1.,  1.,  2.,  3.,  0.,  0.,  1.],
       [ 1.,  0.,  1.,  3.,  0.,  0.,  1.],
       [ 1.,  0.,  3.,  3.,  0.,  0.,  1.],
       [ 1.,  0.,  2.,  2.,  0.,  0.,  1.],
       [ 1.,  0.,  2.,  4.,  0.,  0.,  1.],
       [ 1.,  0.,  2.,  3.,  1.,  0.,  1.],
       [ 1.,  0.,  2.,  3.,  0.,  1.,  1.],
       [ 1.,  0.,  2.,  3.,  0.,  0.,  0.],
       [ 1.,  0.,  2.,  3.,  0.,  0.,  2.]])

I would like to obtain maximally fast implementation of similar_vectors(vector) function mentioned above. It is run in my simulation millions of times for vector of length ~few dozens, so speed is crucial. I am interested both in pure numpy solutions as well as some wrappers to other languages. The code can be parallel.

My current implementation is the following:

def singleOne(length,positionOne): #generates a vector of len length filled with zeros apart from a single one at positionOne
    arr=np.zeros(length)
    arr[positionOne]=1
    return arr

def similar_vectors(state):
    connected=[]
    for i in range(len(state)):
        if(state[i]!=0):
            connected.append(state-singleOne(state.shape[0],i))       
        connected.append(state+singleOne(state.shape[0],i))
    return np.array(connected)

This is terribly slow due to the for loop that I can't easily get rid of.

For reference I'm attaching the profile from 1000 executions of similar_vectors(vector):

         37003 function calls in 0.070 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.070    0.070 {built-in method builtins.exec}
        1    0.003    0.003    0.070    0.070 <string>:2(<module>)
     1000    0.035    0.000    0.064    0.000 <ipython-input-19-69431411f902>:6(similar_vectors)
    11000    0.007    0.000    0.021    0.000 <ipython-input-19-69431411f902>:1(singleOne)
    11000    0.014    0.000    0.014    0.000 {built-in method numpy.core.multiarray.zeros}
     2000    0.009    0.000    0.009    0.000 {built-in method numpy.core.multiarray.array}
    11000    0.002    0.000    0.002    0.000 {method 'append' of 'list' objects}
     1000    0.000    0.000    0.000    0.000 {built-in method builtins.len}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

Answer 1

Here's a vectorized approach. You could create a diagonal matrix of 1s and another of -1s, then add the former to the original array and separately the latter on those positions where the original array is not 0. Then use np.concatenate to concatenate both ndarrays:

def similar_vectors(a):
    ones = np.ones(len(a))
    w = np.flatnonzero(a!=0)
    return np.concatenate([np.diag(-ones)[w]+a, np.diag(ones)+a])

---

Sample run

a = np.array([1,0,2,3,0,0,1])
similar_vectors(a)

array([[0., 0., 2., 3., 0., 0., 1.],
       [1., 0., 1., 3., 0., 0., 1.],
       [1., 0., 2., 2., 0., 0., 1.],
       [1., 0., 2., 3., 0., 0., 0.],
       [2., 0., 2., 3., 0., 0., 1.],
       [1., 1., 2., 3., 0., 0., 1.],
       [1., 0., 3., 3., 0., 0., 1.],
       [1., 0., 2., 4., 0., 0., 1.],
       [1., 0., 2., 3., 1., 0., 1.],
       [1., 0., 2., 3., 0., 1., 1.],
       [1., 0., 2., 3., 0., 0., 2.]])

a = np.array([0,0,0])
similar_vectors(a)

array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])