diff --git a/tree.py b/tree.py
index 2ebc628..1c97076 100644
--- a/tree.py
+++ b/tree.py
@@ -33,12 +33,20 @@ def format(node: scHPL.TreeNode):
 
     def format_name(name: List[str]):
         return " & ".join(name)
-
-    for current in node.walk():
-        cur_name = format_name(current.name)
-        nodes.append({"name": cur_name, "id": hash(cur_name)})
-        for child in current.descendants:
+    
+    def traverse(node: scHPL.TreeNode, level: int = 0):
+        cur_name = format_name(node.name)
+        value = {"name": cur_name, "id": hash(cur_name)}
+        if level <= 1:
+            value["group"] = cur_name
+        elif level > 1:
+            value["group"] = format_name(node.ancestor.name)
+        nodes.append(value)
+        for child in node.descendants:
             links.append({"source": hash(cur_name), "target": hash(format_name(child.name))})
+            traverse(child, level + 1)
+
+    traverse(node)
 
     return {"nodes": nodes, "links": links}
 
@@ -51,4 +59,4 @@ def tree_server(input, output, session, _tree: reactive.Value[scHPL.TreeNode]):
     @output
     @render_tree
     def tree():
-        return format(_tree.get()[0])
\ No newline at end of file
+        return format(_tree.get()[0])
diff --git a/tree/tree.js b/tree/tree.js
index 1e4dee1..c10572e 100644
--- a/tree/tree.js
+++ b/tree/tree.js
@@ -11,6 +11,7 @@ if (Shiny) {
                 .backgroundColor('#101020')
                 .linkColor(() => 'rgba(255,255,255,0.2)')
                 .nodeOpacity(0.9)
+                .nodeAutoColorBy('group')
                 .linkDirectionalParticles(2)
                 .linkDirectionalParticleWidth(0.8)
                 .linkDirectionalParticleSpeed(0.006)