Efficient probability tree branching

Question

As an example, I have an array of branches and probabilities that looks like this:

paths = np.array([
[1, 0, 1.0],
[2, 0, 0.4],
[2, 1, 0.6],
[3, 1, 1.0],
[5, 1, 0.25],
[5, 2, 0.5],
[5, 4, 0.25],
[6, 0, 0.7],
[6, 5, 0.2],
[6, 2, 0.1]])

The columns are upper node, lower node, probability. Here's a visual of the nodes:

        6
     /  |  \
    5   0   2
  / | \    / \
 1  2  4  0   1
 |  /\        |
 0 0  1       0
      |
      0

I want to be able to pick a starting node and output an array of the branches and cumulative probabilities, including all the duplicate branches. For example: start_node = 5 should return

array([
[5, 1, 0.25],
[5, 2, 0.5],
[5, 4, 0.25],
[1, 0, 0.25],
[2, 0, 0.2],
[2, 1, 0.3],
[1, 0, 0.3]])

Notice the [1, 0, x] branch is included twice, as it's fed by both the [5, 1, 0.25] branch and the [2, 1, 0.3] branch.

Here's some code I got working but it's far too slow for my application (millions of branches):

def branch(start_node, paths):
  output = paths[paths[:,0]==start_node]
  next_nodes = output

  while True:
    can_go_lower = np.isin(next_nodes[:,1], paths[:,0])

    if ~np.any(can_go_lower): break

    next_nodes_checked = next_nodes[can_go_lower]

    next_nodes = np.empty([0,3])
    for nodes in next_nodes_checked:
      to_append = paths[paths[:,0]==nodes[1]]
      to_append[:,2] *= nodes[2]
      next_nodes = np.append(next_nodes, to_append, axis=0)

    output = np.append(output, next_nodes, axis=0)

  return output

The branches are always higher to lower, therefor getting caught in circles isn't a concern. A way to vectorize the for loop and avoid the appends would be the best optimization, I think.

user_7831421 · Answer

Instead of storing in numpy array lets' store graph in dict.

tree = {k:arr[arr[:, 0] == k] for k in np.unique(arr[:, 0])}

Make as set of nodes which are non-leaf:

non_leaf_nodes = set(np.unique(arr[:, 0]))

Now to find the branch and cumulative probability:

def branch(start_node, tree, non_leaf_nodes):
curr_nodes = [[start_node, start_node, 1.0]]    #(prev_node, starting_node, current_probability)
output = []
while True:
    next_nodes = []
    for _, node, prob in curr_nodes:
        if node not in non_leaf_nodes: continue
        subtree = tree[node]
        to_append = subtree.copy()
        to_append[:, 2] *= prob
        to_append = to_append.tolist()
        output += to_append
        next_nodes += to_append
    curr_nodes = next_nodes 
    if len(curr_nodes) == 0:
        break
return np.array(output)

Output:

>>> branch(5, tree, non_leaf_nodes)

array([
   [5.  , 1.  , 0.25],
   [5.  , 2.  , 0.5 ],
   [5.  , 4.  , 0.25],
   [1.  , 0.  , 0.25],
   [2.  , 0.  , 0.2 ],
   [2.  , 1.  , 0.3 ],
   [1.  , 0.  , 0.3 ]])

I am expecting it to work faster. Let me know.

Efficient probability tree branching

1 Answers