Search code examples
pythonlistnumpybitwise-operators

variable expansion of ones in python list


I have a python list of variable length filled with 0s and 1s.

I want to create a new list where all 1s are expanded by a certain offset.

Examples:

offset = 1

l1 = [0,0,1,0]
l1_new = l[0,1,1,1]

l2 = [1,0,0,0,1,0,1,0,0]
l2_new = [1,1,0,1,1,1,1,1,0]

My solution code is not very fast and also does not use any numpy / vectorization / bitwise operations. But I guess some of those methods should be applicable here.

offset = 1
l_old = [0,0,1,0]
l_new = []
for i,l in enumerate(l_old):
    hit = False
    for o in range(offset+1)[1:]:
        if (i+o<len(l_old) and l_old[i+o]) or (i>0 and l_old[i-o]):
            hit = True
            break
    if hit or l_old[i]:
        l_new.append(1)
    else:
        l_new.append(0)

Hint: The solution should be fast and generic for any list of 0s and 1s and for any offset


Solution

  • Here is a linear (O(n+offset)) time solution:

    import numpy as np
    
    def symm_dil(a,hw):
        aux = np.zeros(a.size+2*hw+1,a.dtype)
        aux[:a.size] = a
        aux[2*hw+1:] -= a
        return np.minimum(aux.cumsum(),1)[hw:-hw-1]
    
    #example
    rng = np.random.default_rng()
    a = rng.integers(0,2,10)
    print(a)
    print(symm_dil(a,2))
    

    Sample output:

    [0 0 0 1 0 0 0 0 0 1]
    [0 1 1 1 1 1 0 1 1 1]