Question

以下代码为我提供了理想的结果，但我想知道是否有更通用的方法将其应用于1D和2D数组，例如使用花哨/高级索引？

import numpy as np

df1 = np.array([-5, 2, 3, 1, 2, 3])
df2 = np.array([[-5, 2, 3, 1, 2, 3], [-5, 2, 2, 1, 2, 3]])

def func(df):
    if df.ndim == 1:
        cum_cf = df.cumsum()
        ny = np.argmax(cum_cf > 0)-1
        n = cum_cf[ny]
        p = df[ny+1]
    else:
        cum_cf = df.cumsum(axis=1)
        ny = np.argmax(cum_cf > 0, axis=1)-1
        rows = np.arange(df.shape[0])
        n = cum_cf[rows, ny]
        p = cum_cf[rows, ny+1]
    value = 1 + ny - (n/p)
    return value

print(func(df1))
>>> 3.0
print(func(df2))
>>> [3. 4.]

所以这更多是关于清理代码，因为我想我可以省略rows语句，并将变量n和p放在if else条件之外。我可以使用省略号吗？基本上，它是对所有行进行子集设置，但一维数组没有行。

Answer 1

根据评论，这是正确的方法吗？

def func(df):
    df = np.atleast_2d(df)
    cum_cf = df.cumsum(axis=1)
    ny = np.argmax(cum_cf > 0, axis=1)-1
    rows = np.arange(df.shape[0])
    n = cum_cf[rows, ny]
    p = cum_cf[rows, ny+1]
    value = 1 + ny - (n/p)
    return value

独立于尺寸的笨拙的花式/高级索引

1 个答案: