dariosky · March 2, 2018 22:38 · Mar 2, 2018
diff --git a/andre_read.py b/andre_read.py
@@ -0,0 +1,30 @@
+def iterelements(tree):
+    """ Tree is a list, generate the element with the logic of the problem
+        element can be a value or a list, when a list recur
+    """
+    for el in tree:
+        if isinstance(el, list):
+            yield el  # when a list, iterate the list first, then the child
+            yield from iterelements(el)
+        else:
+            yield el
+
+
+def get_n_element(pos, tree):
+    for i, element in enumerate(iterelements(tree)):
+        if i == pos:
+            return element
+
+
+class TestKnownGet:
+    def test_simple(self):
+        assert get_n_element(0, [10, 20, [[30, 40], 50], 60]) == 10
+        assert get_n_element(1, [10, 20, [[30, 40], 50], 60]) == 20
+
+    def test_more(self):
+        assert get_n_element(2, [10, 20, [[30, 40], 50], 60]) == [[30, 40], 50]
+        assert get_n_element(3, [10, 20, [[30, 40], 50], 60]) == [30, 40]
+        assert get_n_element(4, [10, 20, [[30, 40], 50], 60]) == 30
+        assert get_n_element(5, [10, 20, [[30, 40], 50], 60]) == 40
+        assert get_n_element(6, [10, 20, [[30, 40], 50], 60]) == 50
+        assert get_n_element(7, [10, 20, [[30, 40], 50], 60]) == 60
diff --git a/andre_write.py b/andre_write.py
@@ -0,0 +1,42 @@
+# change way - pure
+
+
+def change_n_element_with_pos(pos, tree, newval):
+    """ Same logic as the reader, but doesn't use yield
+        we update the pos as we go, recurring & letting to the recursion do their job
+    """
+    new_tree = []
+    i = 0
+    for element in tree:
+        if i == pos:
+            new_tree.append(newval)
+        else:
+            if isinstance(element, list):
+                new_child, new_pos = change_n_element_with_pos(pos - i - 1, element, newval)
+                new_tree.append(new_child)
+                i += new_pos
+            else:
+                new_tree.append(element)
+        i += 1
+    return new_tree, i
+
+
+def change_n_element(pos, tree, new_val):
+    new_tree, pos = change_n_element_with_pos(pos, tree, new_val)
+    return new_tree
+
+
+class TestKnownChange:
+    def test_change_simple(self):
+        assert change_n_element(1, [10, 20, 30], 0) == [10, 0, 30]
+
+    def test_change_folded(self):
+        assert change_n_element(3, [10, [20, 30], 40], 0) == [10, [20, 0], 40]
+
+    def test_change(self):
+        assert change_n_element(0, [10, 20, [[30, 40], 50], 60], 0) == [0, 20, [[30, 40], 50], 60]
+        assert change_n_element(2, [10, 20, [[30, 40], 50], 60], 0) == [10, 20, 0, 60]
+        assert change_n_element(5, [10, 20, [[30, 40], 50], 60], 0) == [10, 20, [[30, 0], 50], 60]
+
+    def test_post(self):
+        assert change_n_element(7, [10, 20, [[30, 40], 50], 60], [[0]]) == [10, 20, [[30, 40], 50], [[0]]]