From 6c920b34294b7b1e1c04284fd6e78f32e3a86602 Mon Sep 17 00:00:00 2001 From: Philipp Schlegel Date: Sat, 2 Dec 2023 12:04:07 +0000 Subject: [PATCH] add utility function for generating parent indices --- fastcore/utils.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 fastcore/utils.py diff --git a/fastcore/utils.py b/fastcore/utils.py new file mode 100644 index 0000000..47e4c22 --- /dev/null +++ b/fastcore/utils.py @@ -0,0 +1,15 @@ +import numpy as np + + +def node_indices(A, B): + """For each node ID in A find its index in B. + + Typically `A` will be parent IDs and `B` will be node IDs. + Negative IDs (= parents of root nodes) will be passed through. + + Note that there is no check whether all IDs in A actually exist in B. If + an ID in A does not exist in B it gets a negative index (i.e. like roots). + """ + ix_dict = dict(zip(B, np.arange(len(B)))) + + return np.array([ix_dict.get(p, -1) for p in A])