python - 删除 Pandas 后如何更新分层索引？

Question

我像这样设置我的 DataFrame:

cols = ['molecule_id', 'atom_id', 'atom_type', 'x', 'y', 'z']
data = pd.DataFrame()\
         .from_dict(data_dict, orient='index', columns=cols)\
         .set_index(['molceule_id', 'atom_id'])
print(data.head(8))

这里 data_dict: dict[str, list]结果是:

                    atom_type         x         y         z
molecule_id atom_id
0           0               C -2.893477 -2.893477 -2.893477
            1               S -3.293477 -2.893477 -2.893477
1           0               C -2.893477 -1.736086 -1.736086
            1               S -3.293477 -1.736086 -1.736086
2           0               C -1.736086 -2.893477 -1.736086
            1               S -2.136086 -2.893477 -1.736086
3           0               C -1.736086 -1.736086 -2.893477
            1               S -2.136086 -1.736086 -2.893477

在代码的后面，我需要从这个框架中删除一个分子(比如 #1)，因此我这样做:

data.drop(labels=1, level='molecule_id', axis=0, inplace=True)

                    atom_type         x         y         z
molecule_id atom_id
0           0               C -1.736086 -2.893477 -1.736086
            1               S -2.136086 -2.893477 -1.736086
2           0               C -2.893477 -2.893477 -2.893477
            1               S -3.293477 -2.893477 -2.893477
3           0               C -2.893477 -1.736086 -1.736086
            1               S -3.293477 -1.736086 -1.736086
4           0               C -1.736086 -2.893477  0.578695
            1               S -2.136086 -2.893477  0.578695

此时，我想调整 'molecule_id' 索引以实现所需的输出:

                    atom_type         x         y         z
molecule_id atom_id
0           0               C -1.736086 -2.893477 -1.736086
            1               S -2.136086 -2.893477 -1.736086
1           0               C -2.893477 -2.893477 -2.893477
            1               S -3.293477 -2.893477 -2.893477
2           0               C -2.893477 -1.736086 -1.736086
            1               S -3.293477 -1.736086 -1.736086
3           0               C -1.736086 -2.893477  0.578695
            1               S -2.136086 -2.893477  0.578695

设置索引 Pandas 似乎创建了 FrozenList，因此我无法执行以下操作:

data.index.levels[0] = new_id_level

我想到的解决方案是从头开始重建 MultiIndex 并使用 set_index() 应用于 DataFrame :

atoms_per_molecule = 2
num_molecules = len(data)//atoms_per_molecule
molecule_id = np.repeat(range(num_molecules), atoms_per_molecule)
atom_id = np.tile(range(atoms_per_molecule), num_molecules)
tuples = list(zip(molecule_id, atom_id))
names = ['molecule_id', 'atom_id']
multi_id = pd.MultiIndex.from_tuples(tuples, names=names)
data.set_index(multi_id, inplace=True)

它运行良好，但考虑到我计划执行的下降次数，它似乎不合理地复杂。
因此，我想知道是否还有其他更有效的方法可以做到这一点？
P.S:也许可以使用给定的模式创建某种可重置的索引？

Answer 1

示例数据框:

import pandas as pd

data = pd.DataFrame({
    'a': [0, 0, 0, 1, 2, 2, 3, 3],
    'b': [0, 1, 2, 0, 0, 1, 0, 1],
    'col_1': [3, 14, 15, 92, 65, 35, 89, 79],
    'col_2': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']
})
data = data.set_index(['a', 'b'])
data = data.drop(labels=1, level='a', axis=0, inplace=False)
print(data)

这给出:

     col_1 col_2
a b             
0 0      3     a
  1     14     b
  2     15     c
2 0     65     e
  1     35     f
3 0     89     g
  1     79     h

修改索引:

data.index = data.index.remove_unused_levels()
n = data.index.get_level_values(0).nunique()
data.index = data.index.set_levels(range(n), level=0)

当您从数据框中删除行时，这不会更改基础索引。第一行修改索引，删除数据帧中不再使用的索引级别。第二行计算新索引级别 0 中不同值的数量。最后，第三行用连续整数替换级别 0 值。
生成的数据框如下所示:

     col_1 col_2
a b             
0 0      3     a
  1     14     b
  2     15     c
1 0     65     e
  1     35     f
2 0     89     g
  1     79     h