獨立於尺寸的笨拙的花式/高級索引

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