Python:遍历数据框列，检查存储在数组中的条件值，并将值获取到列表

Question

在论坛中获得一些帮助后，我设法完成了我正在寻找的事情，现在我需要进入下一个阶段。 (长解释在这里: Python Data Frame: cumulative sum of column until condition is reached and return the index ):

我有一个数据框:

In [3]: df
Out[3]: 
   index  Num_Albums  Num_authors
0      0          10            4
1      1           1            5
2      2           4            4
3      3           7         1000
4      4           1           44
5      5           3            8

我添加了一列，其中包含另一列的累计总和。

In [4]: df['cumsum'] = df['Num_Albums'].cumsum()

In [5]: df
Out[5]: 
   index  Num_Albums  Num_authors  cumsum
0      0          10            4      10
1      1           1            5      11
2      2           4            4      15
3      3           7         1000      22
4      4           1           44      23
5      5           3            8      26

然后我将条件应用到 cumsum 列，并提取满足给定容差的条件的行的相应值:

In [18]: tol = 2

In [19]: cond = df.where((df['cumsum']>=15-tol)&(df['cumsum']<=15+tol)).dropna()

In [20]: cond
Out[20]: 
   index  Num_Albums  Num_authors  cumsum
2    2.0         4.0          4.0    15.0

现在，我要做的是用存储在数组中的条件替换示例中的条件 15。检查条件何时满足并且不检索整行，而只检索列 Num_Albums 的值。最后，所有这些检索到的值(每个条件一个)都存储在一个数组或列表中。 来自 matlab，我会做这样的事情(对于这种混合的 matlab/python 语法，我深表歉意):

conditions = np.array([10, 15, 23])
for i=0:len(conditions)
   retrieved_values(i) = df.where((df['cumsum']>=conditions(i)-tol)&(df['cumsum']<=conditions(i)+tol)).dropna()

所以对于上面的数据框，我会得到(tol=0):

retrieved_values = [10, 4, 1]

我想要一个解决方案，让我尽可能保留 .where 函数..

Answer 1

一个快速的方法是利用 NumPy 的广播技术作为 this answer 的扩展。来自链接的同一篇文章，尽管实际上询问了与 DF.where 的使用相关的答案。

广播消除了遍历数组的每个元素的需要，同时它非常高效。

这篇文章的唯一补充是使用 np.argmax 沿每一列(遍历 ↓ 方向)获取第一个 True 实例的索引。

conditions = np.array([10, 15, 23])
tol = 0
num_albums = df.Num_Albums.values
num_albums_cumsum = df.Num_Albums.cumsum().values
slices = np.argmax(np.isclose(num_albums_cumsum[:, None], conditions, atol=tol), axis=0)

检索到的切片:

slices
Out[692]:
array([0, 2, 4], dtype=int64)

对应的数组产生:

num_albums[slices]
Out[693]:
array([10,  4,  1], dtype=int64)

---

如果您仍然喜欢使用 DF.where，这里是另一个使用 list-comprehension 的解决方案 -

[df.where((df['cumsum'] >= cond - tol) & (df['cumsum'] <= cond + tol), -1)['Num_Albums']
   .max() for cond in conditions]
Out[695]:
[10, 4, 1]

不满足给定标准的条件将被替换为 -1。这样做会在最后保留 dtype。