How to split a an array into several smaller arrays based on a given criterium between neighbours

user_12234181 0
Viewed 545

Mathematica has two very useful functions to group an array into a list of smaller arrays based on given criteria: Split[] and SplitBy[] which I need to emulate in Python3 code:

Split[list,test] treats pairs of adjacent elements as identical whenever applying the function "test" to them yields True,

SplitBy[list,f] splits list into sublists consisting of runs of successive elements that give the same value when f is applied.

Thus if

a=[2,3,5,7,11,13,17,19,23,29]

Split[a,(#2-#1 < 4)&] gives:

[[2,3,5,7],[11,13],[17,19],[23],[29]]

and SplitBy[a,(Mod[#,2]==0)&] gives:

[[2],[3,5,7,11,13,17,19,23,29]]

In practice the array to be split might be a 2-dimensional table and the test functions might work on the elements in individual columns.

How can this behaviour coded efficiently in Python3?

wolfram-mathematica arrays python split
2 Answers

Don't have a quick answer for the first part of your question, and lenik already provide a nice zip-based solution, but SplitBy can easily be reproduced using the groupby function of the itertools module (doc here).

Beware, groupby will insert a separator (~ create a new group) each time the key change. So if you want something like SplitBy, you have to sort it according to the key function first.

In the end, it will give you something like this:

>>> def split_by(l, func):
        groups = []
        sorted_l = sorted(l, key=func)
        for _, g in it.groupby(sorted_l, key=func):
            groups.append(list(g)) 
        return groups

>>> split_by([2,3,5,7,11,13,17,19,23,29], lambda x: x%2)                                                                                                                                                                     
[[2], [3, 5, 7, 11, 13, 17, 19, 23, 29]]

One-liner version using list comprehension:

splited_by = [list(g) for _, g in it.groupby(sorted(l, key=func), key=func)]

Quick timeit benchmark on my old and broken laptop:

  • itertools version

>>> %timeit split_by([2,3,5,7,11,13,17,19,23,29], lambda x: x%2) 8.42 µs ± 92.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

  • zip version

>>> %timeit split_by([2,3,5,7,11,13,17,19,23,29], lambda x: x%2) 10.8 µs ± 53.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

  • try/catch version

>>> %timeit split_by([2,3,5,7,11,13,17,19,23,29], lambda x: x%2) 12.6 µs ± 162 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

user_14983890

Here's some possible solutions:

a) Using python builtin zip

def split(lst, test):
    res = []
    sublst = []

    for x, y in zip(lst, lst[1:]):
        sublst.append(x)
        if not test(x, y):
            res.append(sublst)
            sublst = []

    if len(lst) > 1:
        sublst.append(lst[-1])
        if not test(lst[-2], lst[-1]):
            res.append(sublst)
            sublst = []

    if sublst:
        res.append(sublst)

    return res


def split_by(lst, test):
    return split(lst, lambda x, y: test(x) == test(y))


if __name__ == '__main__':
    a = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
    print(split(a, lambda x, y: y - x < 4))
    print(split_by(a, lambda x: (x % 2) == 0))

b) for-loop with inner try/except:

def split(lst, test):
    res = []
    sublst = []

    for i, x in enumerate(lst):
        try:
            y = lst[i + 1]
            sublst.append(x)
        except IndexError:
            x, y = lst[i - 1], lst[i]
            sublst.append(y)

        if not test(x, y):
            res.append(sublst)
            sublst = []

    if sublst:
        res.append(sublst)

    return res


def split_by(lst, test):
    return split(lst, lambda x, y: test(x) == test(y))


if __name__ == '__main__':
    a = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
    print(split(a, lambda x, y: y - x < 4))
    print(split_by(a, lambda x: (x % 2) == 0))

Output:

[[2, 3, 5, 7], [11, 13], [17, 19], [23], [29]]
[[2], [3, 5, 7, 11, 13, 17, 19, 23, 29]]
user_38093750
Related