tylerl
tylerl

Reputation: 30847

indexing an array

I find myself frequently making indexed lists from flat ones in Python. This is such a common task that I was wondering if there's a standard utility that I should be using for it.

The context is this: given an array, I need to create a dict of smaller arrays using some key for grouping.

e.g:
["Andy","Alice","Bob","Beth","Charlie"] becomes
{"A":["Andy","Alice"],"B":["Bob","Beth"],"C":["Charlie"]}

My solution looks like this:

def make_index(data,key,value=lambda x:x):
    d={}
    for item in data:
        k = key(item)
        v = value(item)
        try: d[k].append(v)
        except KeyError: d[k]=[v]
    return d

It's simple and all, but am I reinventing something that is implemented better elsewhere?

Upvotes: 3

Views: 221

Answers (2)

Lev Levitsky
Lev Levitsky

Reputation: 65791

Not sure why the itertools answer was deleted, but I was writing one myself:

from itertools import groupby
def make_index(data, key = lambda x: x[0]):
    return {key: list(gr) for key, gr in 
        groupby(sorted(data, key=key), key=key)}

In [3]: make_index(["Andy","Alice","Bob","Beth","Charlie"])
Out[3]: {'A': ['Andy', 'Alice'], 'B': ['Bob', 'Beth'], 'C': ['Charlie']}

In [4]: make_index(["Andy","Alice","Bob","Beth","Charlie"], key=lambda x: len(x))
Out[4]: {3: ['Bob'], 4: ['Andy', 'Beth'], 5: ['Alice'], 7: ['Charlie']}

Upvotes: 3

Martijn Pieters
Martijn Pieters

Reputation: 1121524

You can do the same a little simpler with a defaultdict:

from collections import defaultdict

def make_index(data,key,value=lambda x:x):
    d=defaultdict(list)
    for item in data:
        d[key(item)].append(value(item))
    return d

Using a defaultdict is faster than using .setdefault(), which would be the other option.

Upvotes: 5

Related Questions