closure table in pandas

user_6139149 0
Viewed 164

I'd like to create a closure table with pandas. Suppose you have hierarchical data, something like this with given ID's:

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))

So now I would like to give the table parent_unit_id with something looking like this:

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

In this example every child only has one parent, but what if the frame looked like this instead (last J in unit_2 swaped to 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']
    }
)

So that the parent_unit_id for I would be a list [3, 4]

pandas python
2 Answers

Just build the (predecessor) graph and map it onto the unit_name_id column:

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)

Output

     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]

Swapping the I with the J yields:

     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]

Comparing both solutions yields the following:

%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)

The difference is likely to increase for larger DataFrames. The code for comparing the solutions can be found here

user_40015920

The following should do the trick:

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)

Note that this uses pd.Series.apply() which internally iterates over all rows, and thus is slow. If you need a faster solution, let me know as a comment, we could also speed things up using merge and groupby.

user_31155550
Related