closure_table_pandas.py

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', 'I']
    }
)

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


def solution_bstadlbauer(frame, close):
    unit_name_to_id = {
        unit_name: unit_id
        for unit_name, unit_id
        in close[["unit_name", "unit_name_id"]].values
    }

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

    return close


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']}


def solution_danimesejo(frame, close):
    close['parent_unit_id'] = close['unit_name_id'].map(parents(frame, closure))
    return close