How can I obtain a Treelib representation of the below nested list?

user_16426164 0
Viewed 404

I'm trying to convert a nested list into a Treelib representation.

The desired structure of the tree is such that all elements within a list that are not of the type list are at the same hierarchy. A sub-list in the nested list is the child of the element (this has to be an element which is not of type list) immediately preceding it. For instance,

lst = [1,['a','b','c','d',['s','t',['ab','cd',['a','b'],'ef'],'u'],'f']]

would translate to the Treelib representation

from treelib import Node, Tree

tree = Tree()
tree.create_node(1, "root") # We can assume that the list will contain a root node

tree.create_node('a', 'a', parent='root')
tree.create_node('b', 'b', parent='root')
tree.create_node('c', 'c', parent='root')
tree.create_node('d', 'd', parent='root')
tree.create_node('s', 's', parent='d')
tree.create_node('t', 't', parent='d')
tree.create_node('ab', 'ab', parent='t')
tree.create_node('cd', 'cd', parent='t')
tree.create_node('a', 'a1', parent='cd')
tree.create_node('b', 'b1', parent='cd')
tree.create_node('ef', 'ef', parent='t')
tree.create_node('u', 'u', parent='d')
tree.create_node('f', 'f', parent='root')

tree.show()

Desired output:

1
├── a
├── b
├── c
├── d
│   ├── s
│   ├── t
│   │   ├── ab
│   │   ├── cd
│   │   │   ├── a
│   │   │   └── b
│   │   └── ef
│   └── u
└── f

I would guess that this needs some sort of recursive logic to parse the tree and identify the entire hierarchical structure, but I'm not able to come up with the logic for this. How would I write code to generate the Treelib nodes for an arbitrary nested list (which I have manually written here)? Any help would be appreciated.

Edit: A complication here is that entire chunks of the tree can repeat. For instance,

lst = [1,['a','b','c','d',['s','t',['ab','cd',['a','b'],'ef'],'u'],'f','t',['ab','cd',['a','b'],'ef']]]

should translate to the Treelib representation


tree = Tree()
tree.create_node(1, "root") # We can assume that the list will contain a root node

tree.create_node('a', 'a', parent='root')
tree.create_node('b', 'b', parent='root')
tree.create_node('c', 'c', parent='root')
tree.create_node('d', 'd', parent='root')
tree.create_node('s', 's', parent='d')
tree.create_node('t', 't', parent='d')
tree.create_node('ab', 'ab', parent='t')
tree.create_node('cd', 'cd', parent='t')
tree.create_node('a', 'a1', parent='cd')
tree.create_node('b', 'b1', parent='cd')
tree.create_node('ef', 'ef', parent='t')
tree.create_node('u', 'u', parent='d')
tree.create_node('f', 'f', parent='root')
tree.create_node('t', 't1', parent='root')
tree.create_node('ab', 'ab1', parent='t1')
tree.create_node('cd', 'cd1', parent='t1')
tree.create_node('a', 'a2', parent='cd1')
tree.create_node('b', 'b2', parent='cd1')
tree.create_node('ef', 'ef1', parent='t1')


tree.show()

Desired output:

1
├── a
├── b
├── c
├── d
│   ├── s
│   ├── t
│   │   ├── ab
│   │   ├── cd
│   │   │   ├── a
│   │   │   └── b
│   │   └── ef
│   └── u
├── f
└── t
    ├── ab
    ├── cd
    │   ├── a
    │   └── b
    └── ef

Thanks!

tree list python
1 Answers

You can recursively traverse your list, keeping track of the parent:

import treelib, itertools as it, collections as ct
lst = [1,['a','b','c','d',['s','t',['ab','cd',['a','b'],'ef'],'u'],'f']] 
tree = treelib.Tree()
c = ct.defaultdict(lambda :it.count(1))
def build_tree(d, t, p = None):
   last_p = None
   for a, b in it.groupby(d, key=lambda x:not isinstance(x, list)):
      if a:
         for i in b:
            t.create_node(i, (last_p:=(i if (n:=next(c[i])) == 1 else f'{i}{n}')), parent=p)
      else:
         for i in b:
            build_tree(i, t, p = last_p)

build_tree(lst, tree)
tree.show()

Output:

1
├── a
├── b
├── c
├── d
│   ├── s
│   ├── t
│   │   ├── ab
│   │   ├── cd
│   │   │   ├── a
│   │   │   └── b
│   │   └── ef
│   └── u
└── f

Result on second tree:

lst = [1,['a','b','c','d',['s','t',['ab','cd',['a','b'],'ef'],'u'],'f','t',['ab','cd',['a','b'],'ef']]] 
tree = treelib.Tree()
c = ct.defaultdict(lambda :it.count(1))
build_tree(lst, tree)
tree.show()

Output:

1
├── a
├── b
├── c
├── d
│   ├── s
│   ├── t
│   │   ├── ab
│   │   ├── cd
│   │   │   ├── a
│   │   │   └── b
│   │   └── ef
│   └── u
├── f
└── t
    ├── ab
    ├── cd
    │   ├── a
    │   └── b
    └── ef

build_tree without assignment expression:

def build_tree(d, t, p = None):
   last_p = None
   for a, b in it.groupby(d, key=lambda x:not isinstance(x, list)):
      if a:
         for i in b:
            n = next(c[i])
            last_p = i if n == 1 else f'{i}{n}'
            t.create_node(i, last_p, parent=p)
      else:
         for i in b:
            build_tree(i, t, p = last_p)

Removing repetitive tree components:

from itertools import zip_longest as zl
from contextlib import suppress
def tree_eq(tree, t1, t2):
   if not hasattr(t1, 'tag') or not hasattr(t2, 'tag'):
      return False
   if t1 is None or t2 is None:
      return False
   return t1.tag == t2.tag and t1._identifier < t2._identifier and \
   all(tree_eq(tree, *i) for i in zl(tree.children(t1._identifier), tree.children(t2._identifier)))

def prune_tree(d, tree):
   yield d
   for i in tree.children(getattr(d, 'root', d._identifier)):
       yield from prune_tree(i, tree)

tag_d = ct.defaultdict(list)
for i in prune_tree(tree, tree):
   if hasattr(i, 'tag') and tree.children(getattr(i, 'root', i._identifier)):
      tag_d[i.tag].append(i)

for a, *b in tag_d.values():
   for i in b:
      with suppress(treelib.exceptions.NodeIDAbsentError):
         if tree_eq(tree, a, i):
             tree.remove_node(i._identifier)

tree.show()

Output:

1
├── a
├── b
├── c
├── d
│   ├── s
│   ├── t
│   │   ├── ab
│   │   ├── cd
│   │   │   ├── a
│   │   │   └── b
│   │   └── ef
│   └── u
└── f
user_73267380
Related