Fix a failed count for numbers with more than 1 digit

1 files changed, 8 insertions, 5 deletions

diff --git a/challenge-151/lubos-kolouch/python/ch-1.py b/challenge-151/lubos-kolouch/python/ch-1.py
index fd34a05784..b30797ece3 100644
--- a/challenge-151/lubos-kolouch/python/ch-1.py
+++ b/challenge-151/lubos-kolouch/python/ch-1.py
@@ -1,21 +1,22 @@
 import re
 
 
-def get_min_depth(input: str):
+def get_min_depth(my_input: str):
     """Calculate the depth"""
     # Input: '1 | 2 3 | 4 5'
 
     # iterate through the layers. If the next layer does not have 2^n items,
     # there must be a leaf node
 
-    input = re.sub(r"\s", "", input)
-    layers = input.split("|")
+    layers = my_input.split("|")
 
     for layer_count, _ in enumerate(layers, 1):
         # if not defined means we are at the last layer
-
         try:
-            items_count = len(layers[layer_count])
+            layers[layer_count] = re.sub(r"^\s+", "", layers[layer_count])
+            layers[layer_count] = re.sub(r"\s+$", "", layers[layer_count])
+
+            items_count = len(re.split(r"\s+", layers[layer_count]))
         except IndexError:
             return layer_count
 
@@ -27,4 +28,6 @@ def get_min_depth(input: str):
 
 assert get_min_depth("1 | 2 3 | 4 5") == 2
 assert get_min_depth("1 | 2 3 | 4 *  * 5 | * 6") == 3
+assert get_min_depth("1 | 20 3 | 4 *  * 5 | * 6") == 3
+
 assert get_min_depth("1 | 2 3") == 2