使用 NumPy 和 Pandas 优化 DataFrame 元素左移操作

本教程详细介绍了如何处理 pandas dataframe 中的缺失值（nan），并将其非缺失元素向左移动，填充至每行的起始位置。文章通过结合使用 numpy 的 `argmin` 和 `roll` 函数，提供了一种高效且专业的解决方案，适用于需要对dataframe进行数据规整和预处理的场景。

在数据分析和处理中，我们经常会遇到包含缺失值（NaN）的数据集。有时，为了后续的分析或模型的输入，我们需要对这些缺失值进行规整，例如将每行中的非缺失元素紧密排列在左侧。本教程将介绍一种利用 Pandas 和 NumPy 库高效实现这一操作的方法，尤其适用于处理具有特定结构的 DataFrame。

问题描述与示例数据

假设我们有一个方形的 Pandas DataFrame，其中包含数值和 NaN 值。我们的目标是将每行中的非 NaN 元素向左移动，使其填充该行的起始位置，而 NaN 值则移动到右侧。需要注意的是，本场景中假设 DataFrame 始终是方形的，且第一行不包含任何 NaN 值。

以下是我们的起始 DataFrame 示例：

import pandas as pd
import numpy as np

# 创建示例 DataFrame
data = {
    'A': [10, np.nan, np.nan, np.nan],
    'B': [20, 32, np.nan, np.nan],
    'C': [100, 45, 759, np.nan],
    'D': [50, 63, 98, 32]
}
df = pd.DataFrame(data)

print("原始 DataFrame:")
print(df)

输出：

原始 DataFrame:
      A     B      C     D
0  10.0  20.0  100.0  50.0
1   NaN  32.0   45.0  63.0
2   NaN   NaN  759.0  98.0
3   NaN   NaN    NaN  32.0

我们期望的输出结果是：

Imagine By Magic Studio

AI图片生成器，用文字制作图片

下载

       A     B      C     D
0   10.0  20.0  100.0  50.0
1   32.0  45.0   63.0   NaN
2  759.0  98.0    NaN   NaN
3   32.0   NaN    NaN   NaN

解决方案：结合 NumPy 的 argmin 和 roll

解决此问题的核心思路是逐行处理 DataFrame。对于每一行，我们需要找到第一个非 NaN 元素的位置，然后将该行向左循环移动，使得这个非 NaN 元素成为行的第一个元素。np.roll 函数非常适合进行循环位移操作，而 np.argmin 则能帮助我们找到第一个 NaN 值的索引，从而确定位移量。

核心原理

1. 识别 NaN 值位置： 对于 DataFrame 的每一行（转换为 NumPy 数组），使用 np.isnan() 函数可以生成一个布尔数组，其中 True 表示 NaN，False 表示非 NaN。
2. 确定位移量： 我们需要将第一个非 NaN 值移动到索引 0。这意味着我们需要找到第一个 NaN 值的索引。np.argmin() 函数返回数组中最小值（或第一个 True 值）的索引。当应用于 np.isnan(row) 的结果时，np.argmin(np.isnan(row)) 将返回该行中第一个 NaN 值的索引。这个索引值就是我们需要向左移动的距离。
3. 执行循环位移： np.roll(array, shift) 函数可以对数组进行循环位移。如果 shift 为负数，则表示向左位移。因此，我们将使用 -np.argmin(np.isnan(row)) 作为位移量。

代码实现

以下是实现上述逻辑的 Python 代码：

import pandas as pd
import numpy as np

# 创建示例 DataFrame
data = {
    'A': [10, np.nan, np.nan, np.nan],
    'B': [20, 32, np.nan, np.nan],
    'C': [100, 45, 759, np.nan],
    'D': [50, 63, 98, 32]
}
df = pd.DataFrame(data)

# 应用解决方案
shifted_df = pd.DataFrame([np.roll(row, -np.argmin(np.isnan(row))) for row in df.values],
                          columns=df.columns)

print("\n处理后的 DataFrame:")
print(shifted_df)

代码解析

• df.values: 将 DataFrame 转换为底层的 NumPy 数组，以便进行高效的行级操作。
• for row in df.values: 遍历 NumPy 数组的每一行。
• np.isnan(row): 对于当前行 row，生成一个布尔数组，指示每个元素是否为 NaN。例如，对于 [NaN, 32, 45, 63]，它将生成 [True, False, False, False]。
• np.argmin(np.isnan(row)): 找到布尔数组中第一个 True 的索引。在这个例子中，索引是 0。如果行是 [NaN, NaN, 759, 98]，则索引是 0。如果行是 [32, 45, 63, NaN]，则 np.isnan(row) 是 [False, False, False, True]，np.argmin 将返回 3 (因为 False < True，所以它会找第一个 False 的索引，但我们是找第一个 NaN 的索引，即第一个 True 的索引。这里需要澄清一下 np.argmin 的行为：它返回数组中最小值的索引。在布尔数组中，False 被视为 0，True 被视为 1。所以 np.argmin([True, False, False, False]) 会返回 1 (即 False 的索引)。这与我们期望的“第一个非 NaN 的索引”是相反的。
  • 修正理解： 我们的目标是找到第一个非 NaN 元素，然后将它移动到最左边。如果一行是 [NaN, 32, 45, 63]，第一个非 NaN 是 32，它的索引是 1。我们需要向左移动 1 位。
  • 如果一行是 [NaN, NaN, 759, 98]，第一个非 NaN 是 759，它的索引是 2。我们需要向左移动 2 位。
  • 所以，我们需要的是第一个 非 NaN 元素的索引。
  • np.argmin(~np.isnan(row)) 能够找到第一个 False 在 np.isnan(row) 中的索引，也就是第一个 True 在 ~np.isnan(row) 中的索引。~np.isnan(row) 会将 [True, False, False, False] 变为 [False, True, True, True]。np.argmin([False, True, True, True]) 返回 0。
  • 这个逻辑是正确的：np.argmin(np.isnan(row)) 确实返回的是第一个 True (即 NaN) 的索引。如果一行是 [NaN, 32, 45, 63]，np.isnan(row) 是 [True, False, False, False]，np.argmin 返回 0。这意味着第一个 NaN 在索引 0。我们需要将非 NaN 元素向左移动。
  • 重新思考 np.argmin 的作用： 如果 np.argmin(np.isnan(row)) 返回 k，这意味着从索引 0 到 k-1 都是 NaN，而索引 k 是第一个 NaN。这实际上是我们需要向左移动的量，以便将 k 处的非 NaN 元素（如果存在）移动到最左侧。
  • 更准确的解释： np.argmin(np.isnan(row)) 找到的是第一个 NaN 值的索引。假设这个索引是 k。这意味着从 0 到 k-1 的位置上，要么没有元素（如果行是空的，但这里不是），要么是非 NaN 值。但是根据问题描述，只有第一行没有 NaN，其他行可能从 NaN 开始。
  • 示例分析：
    • 对于 row = [10, 20, 100, 50]：np.isnan(row) -> [F, F, F, F]。np.argmin([F, F, F, F]) -> 0。np.roll(row, -0) -> [10, 20, 100, 50]。正确。
    • 对于 row = [NaN, 32, 45, 63]：np.isnan(row) -> [T, F, F, F]。np.argmin([T, F, F, F]) -> 1 (因为 False < True，argmin 找最小值)。这里是关键点，np.argmin 找的是最小值索引。在布尔数组中，False 是 0，True 是 1。所以它会找到第一个 False 的索引。
    • 如果 np.isnan(row) 是 [T, F, F, F]，那么 np.argmin 返回 1。这意味着第一个 NaN 在索引 0，而第一个非 NaN 在索引 1。我们需要将索引 1 处的元素移动到索引 0。所以位移量应该是 1。
    • 如果 np.isnan(row) 是 [T, T, F, F]，那么 np.argmin 返回 2。这意味着第一个 NaN 在索引 0，第二个 NaN 在索引 1，而第一个非 NaN 在索引 2。我们需要将索引 2 处的元素移动到索引 0。位移量应该是 2。
    • 结论：np.argmin(np.isnan(row)) 实际上返回的是第一个 False (即第一个非 NaN 元素) 的索引。 这正是我们需要的向左位移量。
• np.roll(row, -shift_amount): 对当前行 row 进行循环位移。负号表示向左位移。例如，如果 row 是 [NaN, 32, 45, 63]，shift_amount 是 1。那么 np.roll(row, -1) 将 [NaN, 32, 45, 63] 变为 [32, 45, 63, NaN]。完美符合预期。
• pd.DataFrame(...): 将处理后的行列表重新组合成一个新的 DataFrame，并保留原始列名。

注意事项与总结

1. 性能： 这种逐行迭代并应用 NumPy 函数的方法对于中等大小的 DataFrame 来说是高效的。然而，对于非常庞大的 DataFrame，如果性能成为瓶颈，可能需要探索更底层的 Pandas 或 NumPy 矢量化操作（例如使用 apply 配合 axis=1，但本质上仍是迭代，或者尝试自定义 C/Cython 函数）。但对于此特定问题，当前方案已足够优化。
2. 数据类型： 由于 NaN 值的存在，DataFrame 的数据类型通常会被提升为浮点型（float64）。处理后，如果所有 NaN 都被移到右侧且不再需要，且原始数据是整数，你可能需要手动转换回整数类型（例如 df.astype(int)），但这会再次引入 NaN 无法表示为整数的问题，所以通常保持浮点型是更安全的做法。
3. 通用性： 这种方法不仅限于 NaN 值。你可以修改 np.isnan() 部分来匹配任何你想要“移动”到右侧的特定值，例如 row == 0 或 row == -1。
4. 方形 DataFrame 和第一行无 NaN 的假设： 尽管代码在一般情况下也能工作，但问题中强调的这些假设简化了对边缘情况的考虑。例如，如果第一行也有 NaN，代码依然能正确处理。

通过结合 np.argmin 和 np.roll，我们能够优雅且高效地解决 DataFrame 中非 NaN 元素向左规整的问题，为后续的数据处理步骤奠定基础。