Square of Numbers in Nested List Python 3

Question

The problem I have to solve is one that takes a nested list as an input, and returns the same nested list, except each element is the square of the element that previously existed in that spot. This is my code

>>> def treemap(lst):
...     for element in lst:
...         if element == type(list):
...             return treemap(element)
...         else:
...             element=element**2
...     return lst
>>> lst = [1, 2, 3, [4, [5, 6], 7]]
>>> print(treemap(lst))

Now I am getting an error that is saying 'int' object is not iterable. I'm assuming that means that it's trying to run the loop for an integer type, which doesn't make sense to me, given that I only rerun the function for those of list types.

Answer 1

1. Do not return in the recursive call, you will cease processing all remaining elements once you return

2. element == type(list) is incorrect, because type(list) is <class 'type'> which will never be equal to any item in your list. Use isinstance instead

3. In the base case, you'll need to access the element by index to have changes reflected

---

def treemap(lst):
    for i, element in enumerate(lst):
        if isinstance(element, list):
            treemap(element)
        else:
            lst[i] = lst[i]**2
    return lst

Output:

[1, 4, 9, [16, [25, 36], 49]]

Answer 2

This solution generates a new list using a recursive, ternary list comprehension that recurses on itself if the item n is iterable, otherwise it returns its square.

def square_list(lst):
    return [square_list(n) if hasattr(n, '__iter__') else n ** 2 for n in lst]

>>> square_list(lst)
[1, 4, 9, [16, [25, 36], 49]]

EDIT

It is a ternary list comprehension:

[a if condition(x) else b for x in some_iterable]

# Where condition(x) returns True if condition with argument `x` is True, otherwise False.

Conditional list comprehension:

[x for x in some_iterable if condition]

Answer 3

You need to use isinstance() to check for type, and if the element is a list instead of returning treemap(element) you can assign a[i] to treemap(element) which will run recursively until all the elements are processed. For example:

def treemap(lst):
    for i, element in enumerate(lst):
        if isinstance(element, list):
            lst[i] = treemap(element)
        else:
            lst[i] = element ** 2
    return lst

lst=[1 ,2 , 3, [ 4, [ 5, 6 ], 7 ] ]
print(treemap(lst))

output:

[1, 4, 9, [16, [25, 36], 49]]

Answer 4

Solution using a single list comprehension:

>>> lst = [1, 2, 3, [4, [5, 6], 7]]
>>> [(lambda f, x: f(f, x))(lambda g, x: [g(g, y) for y in x] if isinstance(x, list) else x ** 2, el) for el in lst]
[1, 4, 9, [16, [25, 36], 49]]

Not that I would recommend anyone to use this under normal circumstances.