ndarray element-wise median with completion check

ndarray_median_with_completion.py

'''
$ python ndarray_median_with_completion.py
[[1 2 1]
 [2 1 2]
 [1 2 1]]
'''

from typing import List
import numpy as np

def merge_ndarray_list(a: List[np.ndarray]):
     ans = np.ma.mean([np.ma.masked_array(el, mask= el == 0) for el in a], axis=0)
     return ans

A = np.array([
    [0, 1, 0],
    [1, 0, 1],
    [0, 1, 0]
], dtype=np.int16)

A2 = A*2

B = np.array([
    [1, 0, 1],
    [0, 1, 0],
    [1, 0, 1]
], dtype=np.int16)

Z = np.zeros([3,3], dtype=np.int16)

ndarray_list = [A, A2, A2, B, Z, Z, Z]
ndarray_merged_ma = merge_ndarray_list(ndarray_list)
ndarray_merged = np.round(ndarray_merged_ma.data).astype(np.int16)
print(ndarray_merged)