Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Making a list of indexes into a list of lists

Tags:

python

list

I'm looking for a Python magic method to pack a list of indexes of that sort

[0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 4, 4, 4]

into this, with each index grouped in a specific list :

[[0, 1, 2, 3, 4], [5, 6, 7], [8, 9], [10], [11, 12, 13]]

I have already done it with a list comprehension plus an append loop like the following, but I feel like there's a Python one-liner that could do that. I'm working on lists that sometimes reach 10000+ items, so performance is important.

li = [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 4, 4, 4]

result = [[] for _ in xrange(max(li)+1)]

for i in xrange(len(li)):
    result[li[i]].append(i)
like image 761
Kotch Avatar asked Mar 16 '23 15:03

Kotch


1 Answers

You can use itertools.groupby to group the values. Then calculate the indices based on the lengths of each group, and keep a running count of the starting index for that group.

from itertools import groupby
def index_list(l):
    temp = 0
    index_list = []
    for key, group in groupby(l):
        items = len(list(group))
        index_list.append([i+temp for i in range(items)])
        temp += items
    return index_list

Example

>>> l = [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 4, 4, 4]
>>> index_list(l)
[[0, 1, 2, 3, 4], [5, 6, 7], [8, 9], [10], [11, 12, 13]]
like image 73
Cory Kramer Avatar answered Mar 23 '23 00:03

Cory Kramer