Numpy Fancy / Advanced индексирования независимо от размеров

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.]

Так что это больше о чистке моего кода, так как я предполагаю, что я мог бы опустить инструкцию row и поставить переменные n и p вне условия if else. Могу ли я сделать это с помощью Ellipsis? По сути, это подмножество всех строк, но нет строки для одномерного массива.

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