Pandas => 按组获取第一个和最后一个元素的索引

Question

我有一个数据框，大约有1亿行（内存中占用1.4Gb）

给定输入：

df.head()

Out[1]:
     id    term     x
0     1     A       3
1     1     B       2
2     2     A       1
3     2     B       1
4     2     F       1
5     2     G       1
6     2     Z       1
7     3     K       1
8     3     M       1
9     3     N       1
10    3     Q       1
11    3     R       1
12    3     Z       1
13    4     F       1

我想获取每个id的第一行索引。示例：

Out[1]:
     id    first_idx
0     1    0       
1     2    2       
2     3    7      
2     4    13

我的当前方法非常慢：

first_row = {}
last_id = None
first_row = None

#iterate over all rows
for idx,r in bow.iterrows():
    cid = r['id']
    if cid != last_id: #is this an ID we haven't seen before?
        first_row[cid] = idx
        last_id = cid

任何建议都将是巨大的帮助。

- Alexander David

---

通常情况下，您应该避免在pandas数据框中迭代行。 - James

Answer 1

一、涉及通用情况

方法 #1 使用 np.unique -

idx = np.unique(df.id.values, return_index=1)[1]

要获取每个 ID 的最后索引，只需使用反转版本，并从数据帧的长度中减去 -

len(df)-np.unique(df.id.values[::-1], return_index=1)[1]-1

二、当 id 列已经排序好

方法二-A 我们可以使用切片操作来提升性能，因为我们可以避免进行排序 -

a = df.id.values
idx = np.concatenate(([0],np.flatnonzero(a[1:] != a[:-1])+1))

方法二-B 采用 掩码（适合大量身份证号码）

a = df.id.values
mask = np.concatenate(([True],a[1:] != a[:-1]))
idx = np.flatnonzero(mask)

获取最后一个索引位置：

np.flatnonzero(np.concatenate((a[1:] != a[:-1],[True])))

方法三 对于连续的数字，我们可以使用np.bincount -

a = df.id.values
idx = np.bincount(a).cumsum()[:-1]

样例运行 -

In [334]: df
Out[334]: 
    id term  x
0    1    A  3
1    1    B  2
2    2    A  1
3    2    B  1
4    2    F  1
5    2    G  1
6    2    Z  1
7    3    K  1
8    3    M  1
9    3    N  1
10   3    Q  1
11   3    R  1
12   3    Z  1
13   4    F  1

In [335]: idx = np.unique(df.id.values, return_index=1)[1]

In [336]: idx
Out[336]: array([ 0,  2,  7, 13])

如果您需要输出到数据框中 -

In [337]: a = df.id.values

In [338]: pd.DataFrame(np.column_stack((a[idx], idx)), columns=[['id','first_idx']])
Out[338]: 
   id  first_idx
0   1          0
1   2          2
2   3          7
3   4         13

- Divakar

---

这真是飞快。我真的需要学习如何使用numpy进行数据操作。谢谢！ - Alexander David

Answer 2

使用 DataFrameGroupBy.agg：

df = df.index.to_series().groupby(df['id']).first().reset_index(name='x')
print (df)
   id   x
0   1   0
1   2   2
2   3   7
3   4  13

如果想要获取最后一个索引值：

df = df.index.to_series().groupby(df['id']).agg(['first','last']).reset_index()
print (df)
   id  first  last
0   1      0     1
1   2      2     6
2   3      7    12
3   4     13    13

- jezrael

---

这个是在内存使用方面最高效的。其他的答案我之前用过，但是对于占用机器最大内存的大型数据框并不起作用。 - Isaac Sim

Answer 3

为了完整起见：

df.reset_index().groupby("id")["index"].first()
id
1     0
2     2
3     7
4    13

关于last：

df.reset_index().groupby("id")["index"].last()
id
1     1
2     6
3    12
4    13

时间：

df.drop_duplicates('id').reset_index()
1.64 ms ± 31.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

df.reset_index().groupby("id")["index"].first()
1.93 ms ± 8.08 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

df.groupby('id', as_index=False)['x'].agg(lambda x: x.index[0])
7.96 ms ± 419 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

df.index.to_series().groupby(df['id']).first().reset_index(name='x')
3.11 ms ± 77.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

- Sergey Bushmanov

---

df = df.index.to_series().groupby(df['id']).first().reset_index(name='x') 这个怎么样？ - jezrael

对于相同的输出 df.index.to_series().groupby(df['id']).first() ;) - jezrael

@jezrael 我没有编造数字) - Sergey Bushmanov

我认为时间不同 - 第一个返回Series，最后两个返回DataFrame ;) - jezrael

请注意，Numpy的“通用解决方案”比Pandas的解决方案快多次（至少对于“小型”数据框）。 - Skippy le Grand Gourou

Answer 4

....drop_duplicates

df.drop_duplicates('id').reset_index()
Out[564]: 
   index  id term  x
0      0   1    A  3
1      2   2    A  1
2      7   3    K  1
3     13   4    F  1

- BENY

---

已为您添加了计时。到目前为止，您是获胜者。) - Sergey Bushmanov

晚安 @Wen！ - piRSquared