直接上代码:
"""
核心函数:
array_1D!=0 :返回一个True/False序列
array.argmax(axis=0):返回沿axis轴的最大元素的索引,当存在多个相等的最大值时,返回第一个最大值的索引
"""
import numpy as np
array_1D = np.array([0,1,0,-1,0])
array_2D = np.array(
[[0, 1, 0, -1, 0],
[0, 0, -1, 0, 1],
[0, 1, -1, 0, 1],
[4, 0, -1, 0, 1],
[7, 16, -1, 0, 1]])
def get_first_non_zero_1D(array_1D):
first_non_zero = array_1D[(array_1D!=0).argmax(axis=0)]
return first_non_zero
"""
备注:以下三个函数完全等价,个人比较喜欢最后一个 get_first_non_zeros_2D_2,因为看着比较舒服,也可以方便地扩展到更多维度。
"""
def get_first_non_zeros_2D(array_2D):
first_non_zeros = np.array([get_first_non_zero_1D(array_2D[i]) for i in range(array_2D.shape[0])])
return first_non_zeros
def get_first_non_zeros_2D_1(array_2D):
first_non_zeros = []
for i in range(array_2D.shape[0]):
arr = array_2D[i,:]
first_non_zero = arr[(arr!=0).argmax(axis=0)]
first_non_zeros.append(first_non_zero)
return np.array(first_non_zeros)
def get_first_non_zeros_2D_2(array_2D):
none_zero_index = (array_2D!=0).argmax(axis=1)
# first_non_zeros = np.array([array_2D[i,none_zero_index[i]] for i in range(array_2D.shape[0])])
first_non_zeros = array_2D[range(array_2D.shape[0]),none_zero_index]
return first_non_zeros
b = get_first_non_zeros_2D(array_2D)
c = get_first_non_zeros_2D_1(array_2D)
d = get_first_non_zeros_2D_2(array_2D)
print(b)
print(c)
print(d)
"""
打印结果:
[ 1 -1 1 4 7]
[ 1 -1 1 4 7]
[ 1 -1 1 4 7]
"""
参考:https://blog.csdn.net/dlhlSC/article/details/88080283