给定该行的顺序，查找每个 pandas DataFrame 行中的第一个和最后一个元素

Question

我有一个pandasDataFrame，列中有A、B、C 和D 中的值，并且想为每一行确定第一个和最后一个非零列。但是所有行的元素顺序并不相同。它由item_0、item_1 和item_2 列确定。

虽然我可以通过对每一行应用一个函数来轻松地做到这一点，但这对于我的DataFrame 来说变得非常慢。有没有一种优雅的、更 Pythonic/pandasy 的方式来做到这一点？

输入：

   A  B  C  D item_0 item_1 item_2
0  1  2  0  0      A      B      C
1  0  1  1  0      A      B      C
2  1  0  1  0      A      B      C
3  0  2  0  0      D      A      B
4  1  1  0  1      D      A      B
5  0  0  0  1      D      A      B

预期输出：

   A  B  C  D item_0 item_1 item_2 first last
0  1  2  0  0      A      B      C     A    B
1  0  1  1  0      A      B      C     B    C
2  1  0  1  0      A      B      C     A    C
3  0  2  0  0      D      A      B     B    B
4  1  1  0  1      D      A      B     D    B
5  0  0  0  1      D      A      B     D    D

更新：这是apply 的当前代码

import pandas as pd


def first_and_last_for_row(row):
    reference_list = row[["item_0", "item_1", "item_2"]].tolist()
    list_to_sort = (
        row[["A", "B", "C", "D"]].index[row[["A", "B", "C", "D"]] > 0].tolist()
    )
    ordered_list = [l for l in reference_list if l in list_to_sort]
    if len(ordered_list) == 0:
        return None, None
    else:
        return ordered_list[0], ordered_list[-1]


df = pd.DataFrame(
    {
        "A": [1, 0, 1, 0, 1, 0],
        "B": [2, 1, 0, 2, 1, 0],
        "C": [0, 1, 1, 0, 0, 0],
        "D": [0, 0, 0, 0, 1, 1],
        "item_0": ["A", "A", "A", "D", "D", "D"],
        "item_1": ["B", "B", "B", "A", "A", "A"],
        "item_2": ["C", "C", "C", "B", "B", "B"],
    }
)

df[["first", "last"]] = df.apply(first_and_last_for_row, axis=1, result_type="expand")

Answer 1

这是一个完全矢量化的 numpy 方法。它不是很复杂，但有很多步骤，所以我还提供了代码的注释版本：

cols = ['A', 'B', 'C', 'D']
a = df[cols].to_numpy()

idx = df.filter(like='item_').replace({k:v for v,k in enumerate(cols)}).to_numpy()
b = a[np.arange(len(a))[:,None], idx] != 0
first = b.argmax(1)
last = b.shape[1]-np.fliplr(b).argmax(1)-1

c = df.filter(like='item_').to_numpy()
df[['first', 'last']] = c[np.arange(len(c))[:,None],
                          np.vstack((first, last)).T]

mask = b[np.arange(len(b)), first]
df[['first', 'last']] = df[['first', 'last']].where(pd.Series(mask, index=df.index))

注释代码：

cols = ['A', 'B', 'C', 'D']

# convert to numpy array
a = df[cols].to_numpy()
# array([[1, 2, 0, 0],
#        [0, 1, 1, 0],
#        [1, 0, 1, 0],
#        [0, 2, 0, 0],
#        [1, 1, 0, 1],
#        [0, 0, 0, 1]])

# get indexer as numpy array
idx = df.filter(like='item_').replace({k:v for v,k in enumerate(cols)}).to_numpy()
# array([[0, 1, 2],
#        [0, 1, 2],
#        [0, 1, 2],
#        [3, 0, 1],
#        [3, 0, 1],
#        [3, 0, 1]])

# reorder columns and get non-zero
b = a[np.arange(len(a))[:,None], idx] != 0
# array([[ True,  True, False],
#        [False,  True,  True],
#        [ True, False,  True],
#        [False, False,  True],
#        [ True,  True,  True],
#        [ True, False, False]])

# first non-zero
first = b.argmax(1)
# array([0, 1, 0, 2, 0, 0])

# last non-zero
last = b.shape[1]-np.fliplr(b).argmax(1)-1
# array([1, 2, 2, 2, 2, 0])

# get back column names from position
c = df.filter(like='item_').to_numpy()
df[['first', 'last']] = c[np.arange(len(c))[:,None],
                          np.vstack((first, last)).T]

# optional
# define a mask in case a zero was selected
mask = b[np.arange(len(b)), first]
# array([ True,  True,  True,  True,  True,  True])
# mask where argmax was 0
df[['first', 'last']] = df[['first', 'last']].where(pd.Series(mask, index=df.index))

输出：

   A  B  C  D item_0 item_1 item_2 first last
0  1  2  0  0      A      B      C     A    B
1  0  1  1  0      A      B      C     B    C
2  1  0  1  0      A      B      C     A    C
3  0  2  0  0      D      A      B     B    B
4  1  1  0  1      D      A      B     D    B
5  0  0  0  1      D      A      B     D    D

Answer 2

让我第一次尝试“优化”，只是避免内部循环。 这里的解决方案在 60k 行上大约快 1.7 倍（我没有耐心等待 600k）

def first_and_last(row):
    
    # select order given by items 
    i0, i1, i2 = items = np.array(row[["item_0", "item_1", "item_2"]])
    
    # select values in right order
    v0, v1, v2 = values = np.array(row[[i0, i1, i2]])
    
    pos_values = (values > 0)
    n_positives = np.sum(values)
    
    if n_positives == 0:
        return np.nan, np.nan
    else:
        return items[pos_values][[0, -1]]

然后：

df_ = pd.concat([df]*10_000)

# Original function
%time df_.apply(first_and_last_for_row, axis=1, result_type="expand")
CPU times: user 53.3 s, sys: 22.5 ms, total: 53.4 s
Wall time: 53.4 s

# New function
%time df_.apply(first_and_last, axis=1, result_type="expand")
CPU times: user 32.9 s, sys: 0 ns, total: 32.9 s
Wall time: 32.9 s

但是，apply 方法并不是最优的，还有其他方法可以迭代数据帧。特别是，您可以使用itertuples 方法：

def first_and_last_iter(row):
    
    # select order given by items 
    i0, i1, i2 = items = np.array([getattr(row, "item_0"), getattr(row, "item_1"),getattr(row, "item_2")])
    
    # select values in right order
    v0, v1, v2 = values = np.array([getattr(row, i0), getattr(row, i1),getattr(row,i2)])
    
    pos_values = (values > 0)
    n_positives = np.sum(values)
    
    if n_positives == 0:
        return np.nan, np.nan
    else:
        return items[pos_values][[0, -1]]

%time df_[["first", "last"]] = [first_and_last_iter(row) for row in df_.itertuples()]
CPU times: user 1.05 s, sys: 0 ns, total: 1.05 s
Wall time: 1.05 s

这是 50 倍的改进

Answer 3

假设你的 DataFrame 被命名为df，这里有一些使用过滤和没有循环的东西。它也适用于全零行（在这种情况下，值将是NaN）。

在我的机器上，它在大约 13 秒内运行 10,000,000 行。

# create filters stating if each column <item_n> is not zero
i0 = df.lookup(df.index, df.item_0).astype(bool)  # [True, False, True, False, True, True]
i1 = df.lookup(df.index, df.item_1).astype(bool)
i2 = df.lookup(df.index, df.item_2).astype(bool)

# for the "first" column, fill with value of item_0 if column is not zero
df['first'] = df.item_0[i0]  # ['A', NaN, 'A', NaN, 'D', 'D']
# fill the Nans with values of item_1 if column is not zero
df['first'][~i0 & i1] = df.item_1[~i0 & i1]
# fill the remaining Nans with values of item_2 if column is not zero
df['first'][~i0 & ~i1 & i2] = df.item_2[~i0 & ~i1 & i2]

# apply the same logic in reverse order for "last"
df['last'] = df.item_2[i2]
df['last'][~i2 & i1] = df.item_1[~i2 & i1]
df['last'][~i2 & ~i1 & i0] = df.item_0[~i2 & ~i1 & i0]

输出：

   A  B  C  D item_0 item_1 item_2 first last
0  1  2  0  0      A      B      C     A    B
1  0  1  1  0      A      B      C     B    C
2  1  0  1  0      A      B      C     A    C
3  0  2  0  0      D      A      B     B    B
4  1  1  0  1      D      A      B     D    B
5  0  0  0  1      D      A      B     D    D

Answer 4

df = pd.DataFrame(
{
    "A": [1, 0, 1, 0, 1, 0],
    "B": [2, 1, 0, 2, 1, 0],
    "C": [0, 1, 1, 0, 0, 0],
    "D": [0, 0, 0, 0, 1, 1],
    "item_0": ["A", "A", "A", "D", "D", "D"],
    "item_1": ["B", "B", "B", "B", "B", "B"],
    "item_2": ["C", "C", "C", "A", "A", "A"],
}

)

first = []
last = []
for i in range(df.shape[0]):
   check1 = []
   for j in df.columns:
       t1 = list(df.loc[i:i][j].values)[0]
       try:
          if t1 > 0:
             check1.append(j)
       except TypeError:
         continue

 if len(check1) == 2:
    first.append(check1[0])
    last.append(check1[1])
    check1.clear()
 elif len(check1) == 3:
    first.append(check1[2])
    last.append(check1[1])
    check1.clear()
 elif len(check1) == 1:
    first.append(check1[0])
    last.append(check1[0])
    check1.clear()

输出：

Answer 5

def function1(ss:pd.Series):
    ss1=ss.loc[ss.iloc[4:].tolist()]
    ld1=lambda ss2:ss2.loc[lambda ss3:(ss3>0).cumsum()==1].head(1).index.values[0]

    return pd.Series([ld1(ss1),ld1(ss1[::-1])],index=['first','last'])

df1.join(df1.apply(function1,axis=1))


  A  B  C  D item_0 item_1 item_2 first last
0  1  2  0  0      A      B      C     A    B
1  0  1  1  0      A      B      C     B    C
2  1  0  1  0      A      B      C     A    C
3  0  2  0  0      D      A      B     B    B
4  1  1  0  1      D      A      B     D    B
5  0  0  0  1      D      A      B     D    D