Pandas分组保留首行，其余变NaN处理

1. 问题背景与挑战

在数据处理和分析中，我们经常遇到需要对DataFrame进行分组操作的场景。一个常见需求是，在每个分组内部，我们可能只关心第一行的某些特定信息，而希望将后续行的这些信息清空（设置为NaN），同时保持其他列的数据不变。

例如，给定以下DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(
    {
        'a': [
            'a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'b',
        ],
        'b': [
            -20, 20, 20, 20,-70, -70,-11, -100, -1, -1, -100, 100
        ],
        'c': [
            'f', 'f', 'f', 'f', 'f', 'x', 'x', 'k', 'k', 'k', 'k', 'k'
        ],
        'x': [
            'p', 'p', 'p', 'p', 'p', 'x', 'x', 'i', 'i', 'i', 'i', 'i'
        ],
    }
)
print("原始DataFrame:")
print(df)

原始DataFrame:

    a    b  c  x
0   a  -20  f  p
1   a   20  f  p
2   a   20  f  p
3   a   20  f  p
4   a  -70  f  p
5   a  -70  x  x
6   b  -11  x  x
7   b -100  k  i
8   b   -1  k  i
9   b   -1  k  i
10  b -100  k  i
11  b  100  k  i

我们的目标是：

1. 以列a进行分组。
2. 对于每个组，保留其第一行中列b和c的值。
3. 将每个组中除了第一行以外的行，其列b和c的值设置为NaN。
4. 列a和x的值应保持不变，不被NaN化。

期望的输出如下：

    a     b    c  x
0   a -20.0    f  p
1   a   NaN  NaN  p
2   a   NaN  NaN  p
3   a   NaN  NaN  p
4   a   NaN  NaN  p
5   a   NaN  NaN  x
6   b -11.0    x  x
7   b   NaN  NaN  i
8   b   NaN  NaN  i
9   b   NaN  NaN  i
10  b   NaN  NaN  i
11  b   NaN  NaN  i

一个常见的初步尝试是使用groupby().apply()结合iloc和get_loc。虽然这种方法对于少量列可行，但当需要处理数百列时，手动指定每一列或在循环中迭代列会变得非常低效且难以维护。

# 低效的尝试（不推荐用于大量列）
def func(g):
    # 假设我们知道要处理的列是'b'和'c'
    g.iloc[1:, g.columns.get_loc('b')] = np.nan
    g.iloc[1:, g.columns.get_loc('c')] = np.nan
    return g

# df_modified = df.groupby('a', as_index=False).apply(func)
# print(df_modified)

这种方法需要为每个需要NaN化的列单独操作，或者在一个循环中完成，这对于大型DataFrame和大量列来说效率不高。

2. 高效的矢量化解决方案

Pandas提供了强大的矢量化操作，可以更高效地解决这类问题。我们将利用df.duplicated()、df.where()和df.fillna()的组合来实现目标。

2.1 核心思路

1. 识别非首行： 使用df['a'].duplicated()来标记每个分组中除第一行之外的所有行。
2. 条件性NaN化： 使用df.where()结合上述标记，将所有非首行的值（除了分组键本身）设置为NaN。
3. 恢复特定列： 使用df.fillna()，根据原始DataFrame中需要保留的列（例如a和x）来填充之前被NaN化的这些列。

2.2 详细步骤与代码实现

import pandas as pd
import numpy as np

# 重新创建原始DataFrame以确保操作的独立性
df = pd.DataFrame(
    {
        'a': [
            'a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'b',
        ],
        'b': [
            -20, 20, 20, 20,-70, -70,-11, -100, -1, -1, -100, 100
        ],
        'c': [
            'f', 'f', 'f', 'f', 'f', 'x', 'x', 'k', 'k', 'k', 'k', 'k'
        ],
        'x': [
            'p', 'p', 'p', 'p', 'p', 'x', 'x', 'i', 'i', 'i', 'i', 'i'
        ],
    }
)

# 步骤1: 识别每个分组中的重复行（即非首行）
# df['a'].duplicated() 会为每个'a'组的第一个出现返回False，后续重复出现返回True
# ~df['a'].duplicated() 则会为每个'a'组的第一个出现返回True，后续重复出现返回False
mask = ~df['a'].duplicated()
# print("重复行掩码 (~df['a'].duplicated()):")
# print(mask)

# 步骤2: 使用df.where()进行条件性替换
# df.where(condition) 会保留condition为True的元素，将condition为False的元素替换为NaN
# 此时，所有非首行的值（包括'a'、'b'、'c'、'x'）都会被替换为NaN
df_temp = df.where(mask)
# print("\n经过df.where(mask)处理后的DataFrame:")
# print(df_temp)

# 步骤3: 使用df.fillna()恢复需要保留原始值的列
# 我们希望保留'a'和'x'列的原始值。
# df.fillna(other_df) 会根据other_df来填充df中的NaN值。
# 只有在df中为NaN且other_df中对应位置有非NaN值时，才会进行填充。
# 并且，填充操作只针对other_df中存在的列进行。
columns_to_preserve = ['a', 'x']
df_final = df_temp.fillna(df[columns_to_preserve])

print("\n最终处理结果:")
print(df_final)

输出结果：

最终处理结果:
    a     b    c  x
0   a -20.0    f  p
1   a   NaN  NaN  p
3   a   NaN  NaN  p
2   a   NaN  NaN  p
4   a   NaN  NaN  p
5   a   NaN  NaN  x
6   b -11.0    x  x
7   b   NaN  NaN  i
8   b   NaN  NaN  i
9   b   NaN  NaN  i
10  b   NaN  NaN  i
11  b   NaN  NaN  i

注意：输出的行索引顺序可能与原始示例略有不同，这是因为Pandas在处理过程中可能会调整索引，但这不影响数据的逻辑对应关系。如果需要严格的索引顺序，可以在操作后进行sort_index()或reset_index()。

2.3 关键概念解析

• df['a'].duplicated():
  • 此方法用于标记DataFrame中a列的重复值。
  • 默认情况下（keep='first'），它会为每个分组中的第一个出现返回False，而为后续的重复出现返回True。
  • 例如，对于a列中的第一个'a'，返回False；对于第二个'a'，返回True，以此类推。
• ~df['a'].duplicated():
  • ~是逻辑非操作符。它将duplicated()的结果取反。
  • 因此，它会为每个分组中的第一个出现返回True，而为后续的重复出现返回False。这个布尔Series就是我们用来识别首行的掩码。
• df.where(condition):
  • where()方法是一个强大的条件选择工具。
  • 它会根据condition（一个布尔Series或DataFrame）来选择数据。
  • 当condition中的值为True时，df.where()保留DataFrame中对应位置的原始值。
  • 当condition中的值为False时，df.where()会将DataFrame中对应位置的值替换为NaN（默认行为）。
  • 在本例中，df.where(~df['a'].duplicated())会保留每个分组的第一行，并将所有后续行的所有列（包括a、b、c、x）都设置为NaN。
• df.fillna(other_df):
  • fillna()用于填充DataFrame中的NaN值。
  • 当other_df是一个DataFrame时，fillna()会尝试根据other_df中对应列和索引的值来填充当前DataFrame中的NaN。
  • 具体来说，它会查找当前DataFrame中为NaN的位置，如果other_df在相同列和索引位置有非NaN值，则用other_df的值进行填充。
  • 在本例中，df_temp在非首行的所有列都变成了NaN。df_temp.fillna(df[['a', 'x']])会查看df_temp中的NaN。对于a列和x列中的NaN，它会用原始df中对应a列和x列的值来填充。而b列和c列的NaN不会被填充，因为df[['a', 'x']]中不包含b和c列。

3. 注意事项与拓展

• 列的选择： columns_to_preserve = ['a', 'x']这一步非常关键。它明确指定了哪些列在非首行时也应该保留其原始值。如果你希望除了分组列之外的所有列在非首行时都变为NaN，那么columns_to_preserve就只包含分组列即可（例如['a']）。
• 性能： 这种基于where和fillna的矢量化方法在处理大型DataFrame时比groupby().apply()结合行迭代的方式效率高得多，因为它利用了Pandas底层的优化C/Cython实现。
• 通用性： 这种模式可以轻松推广到任何分组键和任意数量的需要保留或NaN化的列。
• 数据类型： 当将数值列中的值替换为NaN时，如果该列原本是整数类型，Pandas会自动将其转换为浮点数类型（因为NaN在Pandas中通常表示为浮点数）。例如，b列从int64变为float64。这是预期行为。

4. 总结

通过巧妙地结合duplicated()、where()和fillna()这三个Pandas函数，我们能够高效且灵活地实现DataFrame分组数据的首行保留与其余值NaN化处理。这种方法不仅代码简洁，而且在处理大规模数据集时表现出卓越的性能，是Pandas数据操作中值得掌握的实用技巧。