从应用于（多个）pandas 列的函数创建 numpy 数组

Question

我有 pd.DataFrame 个值行：

import pandas as pd

df = pd.DataFrame({"col1": [1, 2, 3, 4, 5, 6], "col2": [6, 5, 4, 3, 2, 1]})

我现在想找到一种有效的方法来根据应用于两列的函数的输出创建 np.array 矩阵：

def my_function(x1, x2, y1, y2):
    return x1 > y1 and x2 < y2

解决这个问题的天真的 O(N²) 方法如下：

matrix = []
for _, (x1, x2) in df.iterrows():
    row = []
    for _, (y1, y2) in df.iterrows():
        row.append(my_function(x1, x2, y1, y2))
    matrix.append(row)

给我们：

>>> print(np.array(matrix))

array([[False, False, False, False, False, False],
       [ True, False, False, False, False, False],
       [ True,  True, False, False, False, False],
       [ True,  True,  True, False, False, False],
       [ True,  True,  True,  True, False, False],
       [ True,  True,  True,  True,  True, False]])

是否有更有效的方法可以扩展到更多值？

Answer 1

你可以试试np.vectorize

def my_function(x, y):
    x1, x2 = x
    y1, y2 = y
    return x1 > y1 and x2 < y2


arr = df.to_records(index=False)
f_vfunc = np.vectorize(my_function)
r = f_vfunc(arr[:, None], arr)

print(r)

[[False False False False False False]
 [ True False False False False False]
 [ True  True False False False False]
 [ True  True  True False False False]
 [ True  True  True  True False False]
 [ True  True  True  True  True False]]

Answer 2

这里不需要

numpy.vectorize你可以直接轻松写一个向量代码（而且vectorize不提高速度，它作为一个循环）：

a = df['col1'].to_numpy()
b = df['col2'].to_numpy()

matrix = (a[:,None]>a)&(b[:,None]<b)

输出：

array([[False, False, False, False, False, False],
       [ True, False, False, False, False, False],
       [ True,  True, False, False, False, False],
       [ True,  True,  True, False, False, False],
       [ True,  True,  True,  True, False, False],
       [ True,  True,  True,  True,  True, False]])

速度对比：

%%timeit
f_vfunc(arr[:, None], arr)
37.2 µs ± 256 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%%timeit
(a[:,None]>a)&(b[:,None]<b)
2.44 µs ± 84.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

从应用于（多个）pandas 列的函数创建 numpy 数组

Create numpy array from function applied to (multiple) pandas columns

python

numpy

pandas