Search code examples
pythonnumpyvectorizationmedian

Replace (every) element in a list by the median of the nearest neighbors


I have an array A, say :

import numpy as np
A = np.array([1,2,3,4,5,6,7,8])

And I wish to create a new array B by replacing each element in A by the median of its four nearest neighbors, without taking into account the value at the given position... for example :

B[2] = np.median([A[0], A[1], A[3], A[4]]) (=3) 

The thing is that I need to perform this on a gigantic A and I want to optimize times, so I want to avoid for loops or similar. And... I don't care about the result at the edges.

I already tried scipy.ndimage.filters.median_filter but it is not producing the desired output :

import scipy.ndimage
B = scipy.ndimage.filters.median_filter(A,footprint=[1,1,0,1,1],mode='wrap')

which produces B=[7,4,4,5,6,7,6,6], which is clearly not the correct answer. Any idea is welcome.


Solution

  • On way could be using np.roll to shift the number in your array such as:

    A_1 = np.roll(A,1)
    # output: array([8, 1, 2, 3, 4, 5, 6, 7])
    

    And then the same thing with rolling by -2, -1 and 2:

    A_2 = np.roll(A,2)
    A_m1 = np.roll(A,-1)
    A_m2 = np.roll(A,-2)
    

    Now you just need to sum your 4 arrays, as for each index you have the 4 neighbors in one of them:

    B = (A_1 + A_2 + A_m1 + A_m2)/4.
    

    And as you said you don't care about the edges, I think it works for you!

    EDIT: I guess I was focus on the rolling idea that I mixed up mean and median, the median can be calculated by B = np.median([A_1,A_2,A_m1,A_m2],axis=0)