CODEWITHSUNDEEP
pythonlist
Micrasema
Micrasema

Reputation: 41

make new python lists with indexes from items in a single list

I want to take a list, e.g. [0, 1, 0, 1, 2, 2, 3], and make a list of lists of (index + 1) for each of the unique elements. For the above, for example, it would be [[1, 3], [2, 4], [5, 6], [7]].

Right now my solution is the ultimate in clunkiness:

list_1 = [0, 1, 0, 1, 2, 2, 3]
maximum = max(list_1)
master = []
for i in range(maximum + 1):
    temp_list = []
    for j,k in enumerate(list_1):
        if k == i:
            temp_list.append(j + 1)
    master.append(temp_list)
print master

Any thoughts on a more pythonic way to do this would be much appreciated!

Upvotes: 1

Views: 294

Answers (3)

jonrsharpe
jonrsharpe

Reputation: 122032

I would do this in two steps:

  1. Build a map {value: [list, of, indices], ...}:

    index_map = {}
for index, value in enumerate(list_1):
    index_map.setdefault(value, []).append(index+1)

  2. Extract the value lists from the dictionary into your master list:

    master = [index_map.get(index, []) for index in range(max(index_map)+1)]

For your example, this would give:

>>> index_map
{0: [1, 3], 1: [2, 4], 2: [5, 6], 3: [7]}
>>> master
[[1, 3], [2, 4], [5, 6], [7]]

This implementation iterates over the whole list only once (O(n), where n is len(list_1)) whereas the others so far iterate over the whole list once for each unique element (O(n*m), where m is len(set(list_1))). By taking max(d) rather than max(list_1) you only need to iterate over the length of the unique items, which is also more efficient.

The implementation can be slightly simpler if you make d = collections.defaultdict(list).

Upvotes: 3

Delgan
Delgan

Reputation: 19627

master = [[i+1 for i in range(len(list_1)) if list_1[i]==j] for j in range(max(list_1)+1)]

It is just the same that your current code, but it uses list comprehension which is often a quite good pythonic way to solve this kind of problem.

Upvotes: 0

Thane Plummer
Thane Plummer

Reputation: 10218

list_1 = [0, 1, 0, 1, 2, 2, 3]
master = []
for i in range(max(list_1)+1):
    master.append([j+1 for j,k in enumerate(list_1) if k==i])

Upvotes: 0

Related Questions