Skip to content

Instantly share code, notes, and snippets.

@mesejo
Created November 16, 2020 11:25
Show Gist options
  • Save mesejo/5c06b9ce788177389cf5a9b269a6a553 to your computer and use it in GitHub Desktop.
Save mesejo/5c06b9ce788177389cf5a9b269a6a553 to your computer and use it in GitHub Desktop.
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment