CODEWITHSUNDEEP
pythonpython-3.x
AlTs
AlTs

Reputation: 301

Extending list by adding element to special position

I have to extend list by adding one element. Pos indicates where I should insert new element and also which value should be inserted (the value of element in original list) F.e.

    ([1, 2, 3], 0) => [1, 1, 2, 3]
    ([1, 2, 3], 1) => [1, 2, 2, 3]
    ([1, 2, 3], 2) => [1, 2, 3, 3]

If pos is too high, I had to add 0 element to the end.

([1, 2, 3], 9) => [1, 2, 3, 0]

I have done so far:

def extend_list(numbers, pos):
    if pos > len(numbers) - 1:
        numbers.append(0)
        return numbers
    return numbers[:pos] + numbers[pos] + numbers[pos:]

But I get an error: TypeError: can only concatenate list (not "int") to list What I'm doing wrong?

Upvotes: 2

Views: 100

Answers (2)

cbuch1800
cbuch1800

Reputation: 943

You could alternatively use the .insert() method to make the addition to the list, as I have done below. I also used an else statement instead of 2 return statements out of personal preference.

def extend_list(numbers, pos):
    if pos > len(numbers) - 1:
        numbers.append(0)
    else:
        numbers.insert(pos, numbers[pos])
    return numbers

In fact, due to the nature of lists in python as mutable objects, running this function will edit the original list, meaning the return statement here is unnecessary and the code functions fine without it.

def extend_list(numbers, pos):
    if pos > len(numbers) - 1:
        numbers.append(0)
    else:
        numbers.insert(pos, numbers[pos])

lst = [1,2,3]
extend_list(lst, 1)

## will print [1,2,2,3] ##
print(lst)

Upvotes: 1

Ajax1234
Ajax1234

Reputation: 71451

The issue with your code is that numbers[pos] is an integer that you are trying to add to numbers[:pos] and numbers[pos:], which are lists. A simple solution is to replace return numbers[:pos] + numbers[pos] + numbers[pos:] with return numbers[:pos] + [numbers[pos]] + numbers[pos:]

However, a shorter solution would be:

vals = [([1, 2, 3], 0), ([1, 2, 3], 1), ([1, 2, 3], 2)]
new_vals = [[*a[:b], a[b], *a[b:]] for a, b in vals]

Output:

[[1, 1, 2, 3], [1, 2, 2, 3], [1, 2, 3, 3]]

In Python2:

import itertools
new_vals = [list(itertools.chain.from_iterable([[c, c] if c == a[b] else [c] for c in a])) for a, b in vals]

Output:

[[1, 1, 2, 3], [1, 2, 2, 3], [1, 2, 3, 3]]

Upvotes: 2

Related Questions