CODEWITHSUNDEEP
pythonpython-3.xlist
N.W.
N.W.

Reputation: 373

Create sublists of indexes of equal values from list

I'm trying to split a list of integers into sublists of the the indexes of equal integers. So say I have a list:

original_list = [1,2,1,4,4,4,3,4,4,1,4,3,3]

The desired output would be:

indexes : [[0,2,9], [1], [6,11,12], [3,4,5,7,8,10]]
# corresponds to sublists: [[1,1,1] [2], [3,3,3], [4,4,4,4,4,4]]

I can't figure out how to do this though, as most solutions require you to first sort the original list, but in my case, this messes up the indices. Itertools or np.arrays have not helped me for this reason, as they only group sequential equal elements.

Does anyone know of a solution for this problem? I would love to hear!

Upvotes: 2

Views: 1041

Answers (5)

ninesalt
ninesalt

Reputation: 4354

Here is how I would do it with numpy, using the argsort function I linked in the comments.

original = [1,2,1,4,4,4,3,4,4,1,4,3,3]
indexes = []
s = set()

for n in np.argsort(original):
    if original[n] in s:
        indexes[-1].append(n)
    else:
        indexes.append([n])
        s.add(original[n])

print(indexes)

Upvotes: 1

jpp
jpp

Reputation: 164773

You can use collections.defaultdict for a one-pass solution. Then use sorted if you need, as in your desired result, to sort your indices by value.

original_list = [1,2,1,4,4,4,3,4,4,1,4,3,3]

from collections import defaultdict
from operator import itemgetter

dd = defaultdict(list)

for idx, value in enumerate(original_list):
    dd[value].append(idx)

keys, values = zip(*sorted(dd.items(), key=itemgetter(0)))

print(keys, values, sep='\n')

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

For comparison, the values of dd are insertion ordered in Python 3.6+ (officially in 3.7+, as a CPython implementation detail in 3.6):

print(list(dd.values()))

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

Upvotes: 1

Patrick Haugh
Patrick Haugh

Reputation: 61032

Using enumerate and a defaultdict, you can build a mapping of values to their indices with 

from collections import defaultdict

dd = defaultdict(list)
for index, value in enumerate(original_list):
    dd[value].append(index)

print(dd)
# defaultdict(<class 'list'>, {1: [0, 2, 9], 2: [1], 4: [3, 4, 5, 7, 8, 10], 3: [6, 11, 12]})

Upvotes: 1

Xero Smith
Xero Smith

Reputation: 2076

This can be achieved with a list comprehension.

>>> x = [1,2,1,4,4,4,3,4,4,1,4,3,3]
>>> [[i for i in range(len(x)) if x[i]==y] for y in sorted(set(x))]
[[0, 2, 9], [1], [6, 11, 12], [3, 4, 5, 7, 8, 10]]

Upvotes: 0

Ajax1234
Ajax1234

Reputation: 71461

You can use enumerate:

original_list = [1,2,1,4,4,4,3,4,4,1,4,3,3]
groups = {a:[i for i, c in enumerate(original_list) if c == a] for a in set(original_list)}

Output:

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

Upvotes: 1

Related Questions