CODEWITHSUNDEEP
pythonlist
Kotch
Kotch

Reputation: 334

Making a list of indexes into a list of lists

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)

Upvotes: 4

Views: 386

Answers (4)

mkrieger1
mkrieger1

Reputation: 23139

This can be done with the following expression:

>>> li = [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 4, 4, 4]
>>> [[i for i, n in enumerate(li) if n == x] for x in sorted(set(li))]
[[0, 1, 2, 3, 4], [5, 6, 7], [8, 9], [10], [11, 12, 13]]

Upvotes: 2

Rick
Rick

Reputation: 45251

Not sure if this is better than the other answers, but I found it interesting to work it out nonetheless: 

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

from collections import Counter

result = []
last = 0

for k,v in sorted(Counter(li).items()):
    result.append(list(range(last, last + v)))
    last += v

Upvotes: 2

Pynchia
Pynchia

Reputation: 11590

My implementation:

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

prev = li[0]    
for pos, el in enumerate(li):
    if el == prev:
        lparz.append(pos)
    else:
        lout.append(lparz)
        lparz = [pos,]
    prev = el

lout.append(lparz)
print lout

outputs

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

as required.

Upvotes: 0

Cory Kramer
Cory Kramer

Reputation: 117876

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]]

Upvotes: 3

Related Questions