高效方法返回 Pandas 数据框中的第一个和最后一个条目

Question

我正在尝试实现一种更有效的方法来返回等于特定值的 pandas df 的第一个和最后一个项。我将在下面发布我当前的方法，但可能有更有效的方法。

import pandas as pd

d = ({
    'X' :    ['X','Y','X','Z','X'],
    'Y' :    [2,5,3,5,1],
})

df = pd.DataFrame(data=d)

我想要返回在X == X条件下Y的第一个和最后一个元素。

这是我的尝试，但我认为可能有更高效的方法。

df = df[df['X'] == 'X']
df_first = df.drop_duplicates(subset=['X'], keep = 'first')
df_last = df.drop_duplicates(subset=['X'], keep = 'last')
df1 = pd.concat([df_first, df_last])

# my expected output
df1
   X  Y
0  X  2
4  X  1

- jonboy

Answer 1

使用query（或任何选择方法）和iloc，这应该很简单。

df.query('X == "X"').iloc[[0, -1]]

   X  Y
0  X  2
4  X  1

假设 Y 中没有 NaN 值。否则，请使用 dropna 函数来处理：

df.query('X == "X"').dropna(subset=['Y']).iloc[[0, -1]]

   X  Y
0  X  2
4  X  1

---

另一种使用agg的选项，认为这很有趣。如果你的"Y"包含NaNs，则这很有用。

df.loc[df['Y'].where(df['X'] == 'X').agg(
    ['first_valid_index', 'last_valid_index'])]

   X  Y
0  X  2
4  X  1

- cs95