Pandas闭包表

Question

我想用熊猫做一个闭合表。假设您有分层数据，类似这样的数据具有给定的ID：

df = pd.DataFrame(
    {
        'unit_0': ['A','A','A','A','A','A','A','A'],
        'unit_1': ['B','C','C','C','D','D','E','E'],
        'unit_2': ['F','G','G','H','I','I','J','J']
    }
)

units = [col for col in df]

closure = (df[units].melt(var_name='depth')
                    .drop_duplicates()
                    .rename(columns={'value': 'unit_name'}))

closure['unit_name_id'] = range(0, len(closure))

现在我想给表parent_unit_id一个如下的东西：

depth   unit_name   unit_name_id    parent_unit_id                  
unit_0  A           0               
unit_1  B           1               0
unit_1  C           2               0
unit_1  D           3               0
unit_1  E           4               0
unit_2  F           5               1
unit_2  G           6               2
unit_2  H           7               2
unit_2  I           8               3
unit_2  J           9               4

在本例中，每个子对象只有一个父对象，但如果帧看起来像这样（单位_2中的最后一个J变为I），该怎么办

df = pd.DataFrame(
    {
        'unit_0': ['A','A','A','A','A','A','A','A'],
        'unit_1': ['B','C','C','C','D','D','E','E'],
        'unit_2': ['F','G','G','H','I','I','J','I']
    }
)

因此，I的parent_unit_id将是一个列表[3, 4]

Answer 1

下面应该可以做到这一点：

unit_name_to_id = {
    unit_name: unit_id 
    for unit_name, unit_id 
    in closure[["unit_name", "unit_name_id"]].values
}

def get_parents(df, unit_name_to_id, depth, unit_name):
    unit_number = int(depth.split("_")[1])
    parent_unit_number = unit_number - 1
    parent_unit_column = f"unit_{parent_unit_number}"
    if parent_unit_column not in df:
        return []
    parents = df[df[depth] == unit_name][parent_unit_column]
    return parents.map(unit_name_to_id).unique().tolist()


closure["parent_unit_ids"] = closure \
    .apply(lambda row: get_parents(df, unit_name_to_id, row["depth"], row["unit_name"]), axis=1)

注意，这使用了pd.Series.apply()，它在内部迭代所有行，因此速度很慢。如果您需要一个更快的解决方案，请告诉我，我们也可以使用merge和groupby来加快速度

Answer 2

只需构建（前置）图并将其映射到unit_name_id列：

import pandas as pd
from collections import defaultdict

df = pd.DataFrame(
    {
        'unit_0': ['A', 'A', 'A', 'A', 'A', 'A', 'A', 'A'],
        'unit_1': ['B', 'C', 'C', 'C', 'D', 'D', 'E', 'E'],
        'unit_2': ['F', 'G', 'G', 'H', 'I', 'I', 'J', 'J']
    }
)

units = [col for col in df]
closure = (df[units].melt(var_name='depth')
           .drop_duplicates()
           .rename(columns={'value': 'unit_name'}))
closure['unit_name_id'] = range(0, len(closure))


def parents(frame, close):
    predecessors = defaultdict(set)
    lookup = {k: v for k, v in close[['unit_name', 'unit_name_id']].values}
    for row in frame.values:
        for i, node in enumerate(row[1:], 1):
            predecessors[lookup[node]].add(lookup[row[i - 1]])
    return {k: list(predecessors[k]) or [] for k in close['unit_name_id']}


closure['parent_unit_id'] = closure['unit_name_id'].map(parents(df, closure))

print(closure)

输出

     depth unit_name  unit_name_id parent_unit_id
0   unit_0         A             0             []
8   unit_1         B             1            [0]
9   unit_1         C             2            [0]
12  unit_1         D             3            [0]
14  unit_1         E             4            [0]
16  unit_2         F             5            [1]
17  unit_2         G             6            [2]
19  unit_2         H             7            [2]
20  unit_2         I             8            [3]
22  unit_2         J             9            [4]

将I与J交换产生：

     depth unit_name  unit_name_id parent_unit_id
0   unit_0         A             0             []
8   unit_1         B             1            [0]
9   unit_1         C             2            [0]
12  unit_1         D             3            [0]
14  unit_1         E             4            [0]
16  unit_2         F             5            [1]
17  unit_2         G             6            [2]
19  unit_2         H             7            [2]
20  unit_2         I             8         [3, 4]
22  unit_2         J             9            [4]

比较两种解决方案得出以下结果：

%timeit solution_bstadlbauer(df, closure)
9.29 ms ± 498 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit solution_danimesejo(df, closure)
1.28 ms ± 86.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

对于较大的数据帧，差异可能会增加。比较解决方案的代码可以在here中找到