python获取二维矩阵的每一行的第一个非零元素

最新推荐文章于 2022-12-10 11:52:48 发布

原创

最新推荐文章于 2022-12-10 11:52:48 发布 · 1.9k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#第一个非零元素 #numpy #python #第一个非0元素

本文介绍了如何使用Python编写函数，通过argmax找到一维和二维数组中第一个非零值的索引。通过get_first_non_zero_1D和get_first_non_zeros_2D函数，展示了查找过程并给出了实际运行结果。

直接上代码：

"""
核心函数：
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)

de