CODEWITHSUNDEEP
pythonlist-comprehensionslice
Jack Rolph
Jack Rolph

Reputation: 595

Variable Length Slice in List Comprehension

Suppose I have a numpy array of step sizes, N, and a set of variables V of length np.sum(N). For example:

N = np.array([2,3,2])
V = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])

Preferably using list comprehension, how do I slice V such that the result is a list of lists, split by the steps in N?

For example:

foo(V, N)
> [[0.1,0.2], [0.3, 0.4,0.5], [0.6,0.7]]

Upvotes: 0

Views: 390

Answers (4)

Jack Rolph
Jack Rolph

Reputation: 595

I managed to solve the problem using itertools.islice:

    def UnequalDivide(self, iterable, chunks):
        it = iter(iterable)
        return [list(islice(it, c)) for c in chunks]

Upvotes: 1

Mark
Mark

Reputation: 92460

numpy has a split() function that will give you unequal arrays. It wants indices rather than lengths. You can do this with cumsum() and just ignore the last empty value (or leftovers if you don't account for the whole list).

N = np.array([2,3,2])
V = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])

np.split(V, N.cumsum())[:-1]
# [array([0.1, 0.2]), array([0.3, 0.4, 0.5]), array([0.6, 0.7])]

Upvotes: 1

Paul Panzer
Paul Panzer

Reputation: 53079

How about this?

Vi = iter(V)
[[next(Vi) for _ in range(n)] for n in N]
# [[0.1, 0.2], [0.3, 0.4, 0.5], [0.6, 0.7]]

Upvotes: 1

Red
Red

Reputation: 27577

Here is how you can do that with regular lists and slices:

N = [2,3,2]
V = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]

def foo(N,V):
    n, lst = 0, []
    for i,v in enumerate(N, 1):
        lst.append(V[n:n+v])
        n += v
    return lst

print(foo(N,V))

Output:

[[0.1, 0.2], [0.3, 0.4, 0.5], [0.6, 0.7]]

You can do the same with numpy arrays:

import numpy as np

N = np.array([2,3,2])
V = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])

def foo(N,V):
    n, lst = 0, []
    for i,v in enumerate(N, 1):
        lst.append(V[n:n+v].tolist())
        n += v
    return lst

print(foo(N,V))

Output:

[[0.1, 0.2], [0.3, 0.4, 0.5], [0.6, 0.7]]

Upvotes: 0

Related Questions