CODEWITHSUNDEEP
pythonrecursionnested-lists
DrJessop
DrJessop

Reputation: 472

Square of Numbers in Nested List Python 3

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.

Upvotes: 6

Views: 2827

Answers (7)

englealuze
englealuze

Reputation: 1663

def treemap(f, tree):
    def branchmap(branch):
        # map on branches
        # if it is a branch (a list), treat it as a smaller tree
        # else it is a leave (a number), apply the function on the leave
        return treemap(f, branch) if type(branch) is list else f(branch)
    return list(map(branchmap, tree))

def square(x):
    return x*x

def square_tree(tree):
    return treemap(square, tree)

square_tree([1,[2,3,[4, 5, 6], 7],[8, 9]])
-->>
[1, [4, 9, [16, 25, 36], 49], [64, 81]]

Upvotes: 1

Alexander
Alexander

Reputation: 109546

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]

Upvotes: 2

wllbll
wllbll

Reputation: 551

there are answers provide solutions that in_place change the lst, i'll provide a solution that return a copy of original list.

def treemap(lst):
    cp = list()
    for element in lst:
        if isinstance(element, list):
            cp.append(treemap(element))
        else:
            cp.append(element**2)
    return cp

Upvotes: 0

cs95
cs95

Reputation: 402483

  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]]

Upvotes: 6

DrJessop
DrJessop

Reputation: 472

So as to stay as close to my code as possible, using everyone's suggestions, here is my answer:

1) 1==type(int) is False, the correct way is type(1)==int

2) Changing the element's value does not change the list, so I would need to refer to the index of lst and change its value

3) I should not return the recursive function, but just call it. A return statement passes a value back to the immediate caller of the current function's call-frame, which is not what I want.

Therefore, my final answer is this,

>>> def treemap(lst):
...     for i in range(len(lst)):
...         if type(lst[i])==int:
...             lst[i]=lst[i]**2
...         elif type(lst[i])==list:
...             treemap(lst[i])
...     return lst
>>> lst = [1, 2, 3, [4, [5, 6], 7]]
>>> print(treemap(lst))
[1, 4, 9, [16, [25, 36], 49]]

Upvotes: 0

Jonas Adler
Jonas Adler

Reputation: 10759

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.

Upvotes: 1

Mohd
Mohd

Reputation: 5613

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]]

Upvotes: 1

Related Questions