Refactor day 6

1 files changed, 47 insertions, 63 deletions

diff --git a/solutions/py/06.py b/solutions/py/06.py
index 660093c..2ca5f26 100644
--- a/solutions/py/06.py
+++ b/solutions/py/06.py
@@ -1,14 +1,6 @@
 from typing import Iterable
 import sys
 
-def flatten(l):
-    # https://stackoverflow.com/a/40857703
-    for v in l:
-        if isinstance(v, Iterable):
-            yield from flatten(v)
-        else:
-            yield v
-
 class Tree(object):
     # https://stackoverflow.com/a/28015122
     def __init__(self, name='root', children=None):
@@ -27,65 +19,57 @@ class Tree(object):
             return self.name == other.name
         return False
     def parents(self):
-        if self.parent == None:
-            return []
-        return [self.parent, self.parent.parents()]
-
-
-if __name__ == "__main__":
-    lines = open("../input/06", "r").readlines()
-
-    trees = {}
-    seen = set()  # seen trees. contains only IDs
-    leafs = {}
+        return [self.parent] + self.parent.parents() if self.parent != None\
+                else []
 
-    for orbit in lines:
-        orbit = orbit.rstrip().split(")")
-        #print(orbit)
-        if not orbit[0] in seen:
-            trees[orbit[0]] = Tree(orbit[0])
-            seen.add(orbit[0])
+lines = open("../input/06", "r").readlines()
 
-    # hook up each child to its parent (since every parent now exists)
-    for orbit in lines:
-        orbit = orbit.rstrip().split(")")
-        if orbit[1] in seen:
-            trees[orbit[0]].add_child(trees[orbit[1]])
-        else:
-            leafs[orbit[1]] = Tree(orbit[1])
-            trees[orbit[0]].add_child(leafs[orbit[1]])
-    for k in trees:
-        pass
-        #print(trees[k], trees[k].children)
+trees = {}
+seen = set()  # seen trees. contains only IDs
+leafs = {}
 
-    trees.update(leafs)
+for orbit in lines:
+    orbit = orbit.rstrip().split(")")
+    #print(orbit)
+    if not orbit[0] in seen:
+        trees[orbit[0]] = Tree(orbit[0])
+        seen.add(orbit[0])
 
-    sum_depths = 0
-    depth = 0
-    depths = {}
-    def count_children(tree, depth):
-        global sum_depths  # fusk-rekursion? ville bara att det skulle fungera
-        sum_depths += depth
-        depths[tree.name] = depth
-        for child in tree.children:
-            count_children(child, depth+1)
-    count_children(trees["COM"], 0) 
-    print("1:", sum_depths)
+# hook up each child to its parent (since every parent now exists)
+for orbit in lines:
+    orbit = orbit.rstrip().split(")")
+    if orbit[1] in seen:
+        trees[orbit[0]].add_child(trees[orbit[1]])
+    else:
+        leafs[orbit[1]] = Tree(orbit[1])
+        trees[orbit[0]].add_child(leafs[orbit[1]])
+for k in trees:
+    pass
+    #print(trees[k], trees[k].children)
 
-    parents = {}
-    for tree in trees:
-        parents[tree] = list(flatten(trees[tree].parents()))
+trees.update(leafs)
 
-    # find common parent and distance to it
-    src = trees["YOU"].parent.name
-    tar = trees["SAN"].parent.name
-    src_dist = 0
-    for src_parent in parents[src]:
-        src_dist += 1
-        tar_dist = 0
-        for tar_parent in parents[tar]:
-            tar_dist += 1
-            if src_parent == tar_parent:
-                print("2:", src_parent.name, src_dist, tar_dist, src_dist + tar_dist)
-                sys.exit()
+sum_depths = 0
+depth = 0
+depths = {}
+def count_children(tree, depth):
+    global sum_depths  # fusk-rekursion? ville bara att det skulle fungera
+    sum_depths += depth
+    depths[tree.name] = depth
+    for child in tree.children:
+        count_children(child, depth+1)
+count_children(trees["COM"], 0) 
+print("1:", sum_depths)
 
+# find common parent and distance to it
+src_parents = trees["YOU"].parent.parents()
+tar_parents = trees["SAN"].parent.parents()
+src_dist = 0
+for src_parent in src_parents:
+    src_dist += 1
+    tar_dist = 0
+    for tar_parent in tar_parents:
+        tar_dist += 1
+        if src_parent == tar_parent:
+            print("2:", src_parent.name, src_dist, tar_dist, src_dist + tar_dist)
+            sys.exit()