How to create a D-ary balanced tree in python using `networkx` python package

Question

I am trying to generate a D-ary balanced tree in python using the networkx package.

    import networkx as nx
    g=nx.Graph()
    D= int(input("enter number of children of a node:"));
    L=int(input("Enter the number of levels:"));

    #variable to store the total number of nodes in the tree.
    tot_node=0;

    for i in range(0,L+1):
          tot_node=tot_node+D**i;

    for N in range(1,tot_node):
          for j in range(N,N+D):
                g.add_edge(N,j);
    
    nx.draw(g);

For this I am getting the following tree for D=2 and L=3. enter image description here

Can someone please point out the error in this code? I want to construct a balanced tree for any general D (the number of branches of a node).

user_10919349 · Answer

I have updated the code again to make sure the general cases work. I hope I have not made this more complicated than necessary, I feel like there must be some simpler implementation, maybe one that relies on recursion.

Anyways, I have produced what I think is an acceptable result. Although it is not your code directly, I believe I have implemented something along the lines of the rudimentary solution you want:

import matplotlib.pyplot as plt
import networkx as nx
from networkx import Graph

#We make a node class to track which node to modify (modify here means add children to.)
class Node:
  def __init__(self, node_id, has_children, not_connected):
      self.node_id = node_id
      self.has_children = has_children
      self.not_connected = not_connected

def get_min_not_connected(nodes_tracker):
    smallest = float('inf')
    for node in nodes_tracker:
        #print(f"Is the node {node.node_id} not connected:  {node.not_connected}")
        if node.node_id < smallest and node.not_connected:
            smallest = node.node_id
    return smallest-1

def construction_step(G, node_id, num_children, nodes_tracker):

    #print(f"The range is {len(nodes_tracker)+1} to  {len(nodes_tracker)+num_children+1}")

    #I am just creating new Node objects to track which connections have been made here. Note how the third parameter of not connected is True.


    nodes_tracker = nodes_tracker + [Node(i,False,True) for i in range(len(nodes_tracker)+1, len(nodes_tracker)+num_children+1)]
    for i in range(1, num_children+1):
        print(f'adding edge relation ({node_id}, {get_min_not_connected(nodes_tracker)+i})')

    #Here I am adding the child nodes to the parent ones.
        G.add_edge(node_id, get_min_not_connected(nodes_tracker)+i)
    
    for i in range(1, num_children+1):
        #print(get_min_not_connected(nodes_tracker))
        nodes_tracker[get_min_not_connected(nodes_tracker)].not_connected = False
    
    return nodes_tracker


#Hardcode inputs for your specific example.
#I am using num_children in place of your D variable.
num_children=3
L=2

G=nx.Graph()

#Create the central (initial) node and setup
total_nodes = 0

#correct formula is like 2^0+2^1+...+2^L
for i in range(0,L):
    total_nodes += num_children**i

print(total_nodes)


nodes_tracker = [Node(1,False,False)]


#Create the actual d-ary graph here.
for i in range(1, total_nodes+1):
    nodes_tracker = construction_step(G, i, num_children, nodes_tracker)

#print(len(nodes_tracker))

nx.draw(G); 
plt.show()

For the output with your parameters D=2, L=3, I got:

To test a more general case, I used D=4, L=2 and I got:

And for fun D=5, L=3:

It works with bigger D and L as well, but the charts naturally look very ugly.

Thanks for your patience with this answer and I hope this helps.

How to create a D-ary balanced tree in python using `networkx` python package

1 Answers