From 2bc43652d0ab632cae30dda8bfc671953e2ad420 Mon Sep 17 00:00:00 2001
From: Matej Focko <mfocko@redhat.com>
Date: Sat, 21 May 2022 13:02:05 +0200
Subject: [PATCH] feat: color edges in red-black tree

Signed-off-by: Matej Focko <mfocko@redhat.com>
---
 ranked_tree.js |  8 +++++++-
 rbt.js         | 12 ++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/ranked_tree.js b/ranked_tree.js
index 8ce15e1..4f6a2c3 100644
--- a/ranked_tree.js
+++ b/ranked_tree.js
@@ -37,10 +37,12 @@ class RankedTree {
 
         edges.forEach((vertices) => {
             let [u, v] = vertices;
+            let diff = nodeDifference(v);
+
             result.push(
                 `\t"Node(${u.value})" -> "Node(${
                     v.value
-                })" [label="${nodeDifference(v)}"];\n`
+                })" [label="${diff}"${this.styleEdge(diff)}];\n`
             );
         });
 
@@ -176,4 +178,8 @@ class RankedTree {
     deleteRebalance(node, parent) {
         throw "not implemented!";
     }
+
+    styleEdge(rankDifference) {
+        return "";
+    }
 }
diff --git a/rbt.js b/rbt.js
index ed659dc..809b005 100644
--- a/rbt.js
+++ b/rbt.js
@@ -158,4 +158,16 @@ class RBTree extends RankedTree {
             }
         }
     }
+
+    styleEdge(rankDifference) {
+        switch (rankDifference) {
+            case 2:
+                return " penwidth=2";
+            case 0:
+                return ', color="red"';
+            case 1:
+            default:
+                return "";
+        }
+    }
 }