Giuseppe P.
Giuseppe P.

Reputation: 21

Inserting elements of one list into another list at different positions in python

Consider the two lists: a=[1,2,3] and b=[10,20,30], and a list of positions pos=[p1,p2,p3] giving the positions that the elements of b should take in the final list of 6 elements given by the union of a and b, where p1 is the position of b[0]=10, p2 is the position of b[1]=20 and p3 is the position of b[2]=30.

What is the best python approach to this problem?

Upvotes: 0

Views: 1061

Answers (2)

Mark
Mark

Reputation: 92440

If you make the indices and values into a dictionary, you can then loop over the range of the combined lengths. If the index is in the dict, use the value, otherwise take the next value from a:

a = [1,2,3]
b = [10,20,30]
pos =[2,0,5]

p_b = dict(zip(pos, b))
it_a = iter(a)

[p_b[i] if i in p_b else next(it_a) for i in range(len(a) + len(b))]

# [20, 1, 10, 2, 3, 30]

You will need to insure that the lengths of the arrays and the positions all make sense. If they don't you can run out of a values which will produce a StopIteration exception.

You use a defaultdict for similar approach, which simplifies the list comprehension at the expense of a slightly more complicated setup:

from collections import defaultdict

a = [1,2,3]
b = [10,20,30]
pos =[4,0,2]
it_a = iter(a)

d = defaultdict(lambda: next(it_a))
d.update(dict(zip(pos, b)))

[d[i] for i in range(len(a) + len(b))]
# [20, 1, 30, 2, 10, 3]

Upvotes: 1

Thierry Lathuille
Thierry Lathuille

Reputation: 24232

You could create the output list by extending it with slices of a and appending the next item of b where needed:

def insert(a, b, positions):
    # reorder b and positions so that positions are in increasing order
    positions, b = zip(*sorted(zip(positions, b))) 

    out = []
    a_idx = 0
    it_b = iter(b)
    for pos in positions:
        slice_length = pos - len(out)
        out.extend(a[a_idx:a_idx + slice_length])
        out.append(next(it_b))
        a_idx += slice_length
    out.extend(a[a_idx:])
    return out

An example:

a=[1,2,3]
b=[10,20,30]
pos=[0, 1, 5]

insert(a, b, pos)
# [10, 20, 1, 2, 3, 30]

pos = [0, 2, 4]
insert(a, b, pos)
# [10, 1, 20, 2, 30, 3]

pos=[5, 3, 0]
insert(a, b, pos)
# [30, 1, 2, 20, 3, 10]

Upvotes: 1

Related Questions