python - 使用矢量化在 nan 之后在 numpy 数组的行中获取第一个/第二个/第三个...值

问题描述

我有以下内容pandas df：

| Date | GB | US | CA | AU | SG | DE | FR |
| ---- | -- | -- | -- | -- | -- | -- | -- |
| 1    | 25 |    |    |    |    |    |    |
| 2    | 29 |    |    |    |    |    |    |
| 3    | 33 |    |    |    |    |    |    |
| 4    | 31 | 35 |    |    |    |    |    |
| 5    | 30 | 34 |    |    |    |    |    |
| 6    |    | 35 | 34 |    |    |    |    |
| 7    |    | 31 | 26 |    |    |    |    |
| 8    |    | 33 | 25 | 31 |    |    |    |
| 9    |    |    | 26 | 31 |    |    |    |
| 10   |    |    | 27 | 26 | 28 |    |    |
| 11   |    |    | 35 | 25 | 29 |    |    |
| 12   |    |    |    | 33 | 35 | 28 |    |
| 13   |    |    |    | 28 | 25 | 35 |    |
| 14   |    |    |    | 25 | 25 | 28 |    |
| 15   |    |    |    | 25 | 26 | 31 | 25 |
| 16   |    |    |    |    | 26 | 31 | 27 |
| 17   |    |    |    |    | 34 | 29 | 25 |
| 18   |    |    |    |    | 28 | 29 | 31 |
| 19   |    |    |    |    |    | 34 | 26 |
| 20   |    |    |    |    |    | 28 | 30 |

我已经部分完成了我在这里尝试单独使用 Pandas 做的事情，但这个过程需要很长时间，所以我不得不使用numpy（请参阅在 pandas 列中获取最左边的值），这就是我苦苦挣扎的地方。

本质上，我希望我的函数f接受一个参数int(offset)，从左边捕获nan每一行的第一个非值，并将整个事物作为numpy数组/向量返回，以便：

f(offset=0)


| 0  | 1  |
| -- | -- |
| 1  | 25 |
| 2  | 29 |
| 3  | 33 |
| 4  | 31 |
| 5  | 30 |
| 6  | 35 |
| 7  | 31 |
| 8  | 33 |
| 9  | 26 |
| 10 | 27 |
| 11 | 35 |
| 12 | 33 |
| 13 | 28 |
| 14 | 25 |
| 15 | 25 |
| 16 | 26 |
| 17 | 34 |
| 18 | 28 |
| 19 | 34 |
| 20 | 28 |

正如我在另一篇文章中所描述的，最好想象为每一行从左侧绘制一条水平线，并将与该线相交的值作为数组返回。offset=0然后返回第一个值（在该数组中），offset=1并将返回相交的第二个值，依此类推。

所以：

f(offset=1)

| 0  | 1   |
| -- | --- |
| 1  | nan |
| 2  | nan |
| 3  | nan |
| 4  | 35  |
| 5  | 34  |
| 6  | 34  |
| 7  | 26  |
| 8  | 25  |
| 9  | 31  |
| 10 | 26  |
| 11 | 25  |
| 12 | 35  |
| 13 | 25  |
| 14 | 25  |
| 15 | 26  |
| 16 | 31  |
| 17 | 29  |
| 18 | 29  |
| 19 | 26  |
| 20 | 30  |

上面帖子中提出的pandas解决方案非常有效：

def f(df, offset=0):
    x = df.iloc[:, 0:].apply(lambda x: sorted(x, key=pd.isna)[offset], axis=1)
    return x

print(f(df, 1))

但是，对于较大的迭代，这非常慢。我已经尝试过了，np.apply_along_axis而且速度更慢！

有没有更胖的numpy矢量化方式？

非常感谢。

标签： pythonpandasdataframenumpy

解决方案

麻木的方法

我们可以定义一个函数first_value，它接受一个2D数组和offset(n) 作为输入参数并返回1D数组。基本上，对于每一行，它返回第nth一个值之后的non-nan值

def first_valid(arr, offset=0):
    m = ~np.isnan(arr)
    i =  m.argmax(axis=1) + offset
    iy = np.clip(i, 0, arr.shape[1] - 1)

    vals = arr[np.r_[:arr.shape[0]], iy]
    vals[(~m.any(1)) | (i >= arr.shape[1])] = np.nan
    return vals

熊猫进场

我们可以stack对数据框进行重塑，然后将数据框分组level=0并使用聚合nth，然后reindex根据原始帧来符合聚合帧的索引

def first_valid(df, offset=0):
    return df.stack().groupby(level=0)\
                     .nth(offset).reindex(df.index)

样品运行

>>> first_valid(df, 0)
Date
1     25.0
2     29.0
3     33.0
4     31.0
5     30.0
6     35.0
7     31.0
8     33.0
9     26.0
10    27.0
11    35.0
12    33.0
13    28.0
14    25.0
15    25.0
16    26.0
17    34.0
18    28.0
19    34.0
20    28.0
dtype: float64


>>> first_valid(df, 1)
Date
1      NaN
2      NaN
3      NaN
4     35.0
5     34.0
6     34.0
7     26.0
8     25.0
9     31.0
10    26.0
11    25.0
12    35.0
13    25.0
14    25.0
15    26.0
16    31.0
17    29.0
18    29.0
19    26.0
20    30.0
dtype: float64

>>> first_valid(df, 2)
Date
1      NaN
2      NaN
3      NaN
4      NaN
5      NaN
6      NaN
7      NaN
8     31.0
9      NaN
10    28.0
11    29.0
12    28.0
13    35.0
14    28.0
15    31.0
16    27.0
17    25.0
18    31.0
19     NaN
20     NaN
dtype: float64

表现

# Sample dataframe for testing purpose
df_test = pd.concat([df] * 10000, ignore_index=True)

%%timeit # Numpy approach
_ = first_valid(df_test.to_numpy(), 1)
# 6.9 ms ± 212 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)


%%timeit # Pandas approach
_ = first_valid(df_test, 1)
# 90 ms ± 867 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)


%%timeit # OP's approach
_ = f(df_test, 1)
# 2.03 s ± 183 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

基于 Numpy 的方法大约300x比OP's给定的方法快，而基于 pandas 的方法大约22x更快