Home > Software design >  Python JSON to tree with anytree
Python JSON to tree with anytree

Time:02-03

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?

CodePudding user response:

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))
  • Related