Combine lists having a specific merge order in a pythonic way?

Viewed 1858

I would like to construct list x from two lists y and z. I want all elements from y be placed where ypos elements point. For example:

y = [11, 13, 15]
z = [12, 14]
ypos = [1, 3, 5]

So, x must be [11, 12, 13, 14, 15]

Another example:

y = [77]
z = [35, 58, 74]
ypos = [3]

So, x must be [35, 58, 77, 74]

I've written function that does what I want but it looks ugly:

def func(y, z, ypos):
    x = [0] * (len(y) + len(z))
    zpos = list(range(len(y) + len(z)))
    for i, j in zip(y, ypos):
        x[j-1] = i
        zpos.remove(j-1)
    for i, j in zip(z, zpos):
        x[j] = i
    return x

How to write it in pythonic way?

6 Answers

If the lists are very long, repeatedly calling insert might not be very efficient. Alternatively, you could create two iterators from the lists and construct a list by getting the next element from either of the iterators depending on whether the current index is in ypos (or a set thereof):

>>> ity = iter(y)
>>> itz = iter(z)
>>> syp = set(ypos)
>>> [next(ity if i+1 in syp else itz) for i in range(len(y)+len(z))]
[11, 12, 13, 14, 15]

Note: this will insert the elements from y in the order they appear in y itself, i.e. the first element of y is inserted at the lowest index in ypos, not necessarily at the first index in ypos. If the elements of y should be inserted at the index of the corresponding element of ypos, then either ypos has to be in ascending order (i.e. the first index of ypos is also the lowest), or the iterator of y has to be sorted by the same order as the indices in ypos (afterwards, ypos itself does not have to be sorted, as we are turning it into a set anyway).

>>> ypos = [5,3,1]   # y and z being same as above
>>> ity = iter(e for i, e in sorted(zip(ypos, y)))
>>> [next(ity if i+1 in syp else itz) for i in range(len(y)+len(z))]
[15, 12, 13, 14, 11]

You should use list.insert, this is what it was made for!

def func(y, z, ypos):
    x = z[:]
    for pos, val in zip(ypos, y):
        x.insert(pos-1, val)
    return x

and a test:

>>> func([11, 13, 15], [12, 14], [1,3,5])
[11, 12, 13, 14, 15]

With large lists, it might be a good idea to work with numpy.

Algorithm

  • create a new array as large as y + z
  • calculate coordinates for z values
  • assign y values to x at ypos
  • assign z values to x at zpos

The complexity should be O(n), with n being the total number of values.

import numpy as np

def distribute_values(y_list, z_list, y_pos):
    y = np.array(y_list)
    z = np.array(z_list)
    n = y.size + z.size
    x = np.empty(n, np.int)
    y_indices = np.array(y_pos) - 1
    z_indices = np.setdiff1d(np.arange(n), y_indices, assume_unique=True)
    x[y_indices] = y
    x[z_indices] = z
    return x

print(distribute_values([11, 13, 15], [12, 14], [1, 3, 5]))
# [11 12 13 14 15]
print(distribute_values([77], [35, 58, 74], [3]))
# [35 58 77 74]

As a bonus, it also works fine when ypos isn't sorted:

print(distribute_values([15, 13, 11], [12, 14], [5, 3, 1]))
# [11 12 13 14 15]
print(distribute_values([15, 11, 13], [12, 14], [5, 1, 3]))
# [11 12 13 14 15]

Performance

With n set to 1 million, this approach is a bit faster than @tobias_k's answer and 500 times faster than @Joe_Iddon's answer.

The lists were created this way:

from random import random, randint
N = 1000000
ypos = [i+1 for i in range(N) if random()<0.4]
y = [randint(0, 10000) for _ in ypos]
z = [randint(0, 1000) for _ in range(N - len(y))

Here are the results with %timeit and IPython:

%timeit eric(y, z, ypos)
131 ms ± 1.54 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit tobias(y, z, ypos)
224 ms ± 977 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit joe(y,z, ypos)
54 s ± 1.48 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

Assuming that the ypos indices are sorted, here is another solution using iterators, though this one also supports ypos of unknown or infinite length:

import itertools

def func(y, ypos, z):
    y = iter(y)
    ypos = iter(ypos)
    z = iter(z)
    next_ypos = next(ypos, -1)
    for i in itertools.count(start=1):
        if i == next_ypos:
            yield next(y)
            next_ypos = next(ypos, -1)
        else:
            yield next(z)
Related