Search code examples
pythonjsontree

Python JSON to tree with anytree


I have JSON data that looks something like

{'Tree': [
  {"From": "1",
  "To": "2"},
  {"From": "2",
  "To": "3"}
]}

basically an array of all existing references between nodes. I want to be able to draw trees starting from a chosen root node using this data.

I find that using anytree I must link a node to a parent node, but ideally I want to simply link it to the name and have it put together the tree structure on its own.

# I want to do this
child[0] = Node(data[0]["To"], parent=data[0]["From"])

# But I think I need to assign a node, not its name
child[1] = Node(data[1]["To"], parent=child[?])

Suggestions on how to do this?


Solution

  • Do this in two steps:

    1. Convert the input format into an adjacency list: a dictionary keyed by the node keys, and with associated lists that hold the keys of the connected nodes.

    2. Choose a root node, and then create the anytree tree using the adjacency list.

    Here are two functions that deal with these two tasks:

    from anytree import Node, RenderTree
    from collections import defaultdict 
    
    def makegraph(edges):
        adj = defaultdict(list)
        for edge in edges:
            a, b = edge.values()
            adj[a].append(b)
            adj[b].append(a)
        return adj
    
    def maketree(adj, nodekey):
        visited = set()
        
        def dfs(nodekey, parent):
            if nodekey not in visited:
                visited.add(nodekey)
                node = Node(nodekey, parent)
                for childkey in adj[nodekey]:
                    dfs(childkey, node)
                return node
                    
        return dfs(nodekey, None)
    

    Here is how to use it:

    edges = [
        {"From": "1", "To": "2"},
        {"From": "2", "To": "3"}
    ]
    
    adj = makegraph(edges)
    # Choose "1" as root:
    root = maketree(adj, '1')
    print(RenderTree(root))
    # Now let's choose "2" as root:
    root = maketree(adj, '2')
    print(RenderTree(root))