Hi,
This looks clearer indeed.
One proposal I would have would be to make a big 0-1 matrix based on the tree content. And use that to perform the neighborhoodEmbedding as a single matrix matrix multiplication:
You need a 1 for a node and each of its neighbors and expand it to the feature dimension size.