Itertools groupby 按两个值组织字典列表

Question

I'm attempting to organize values by state of birth as well as if they have 0 money or not. Itertools groupby function looks like the simplest way to do so but I'm struggling implementing it. Open to other options as well.

If I have a list of dictionaries that looks like this

users = [
            {"name": "John", "state_of_birth": "CA", "money": 0},
            {"name": "Andrew", "state_of_birth": "CA", "money": 300},
            {"name": "Scott", "state_of_birth": "OR", "money": 20},
            {"name": "Travis", "state_of_birth": "NY", "money": 0},
            {"name": "Bill", "state_of_birth": "CA", "money": 0},
            {"name": "Mike", "state_of_birth": "NY", "money": 0}
        ]

I'm attempting to get this output

desired_output = [
            [{"name": "John", "state_of_birth": "CA", "money": 0}, {"name": "Bill", "state_of_birth": "CA", "money": 0}],
            [{"name": "Andrew", "state_of_birth": "CA", "money": 300}],
            [{"name": "Scott", "state_of_birth": "OR", "money": 20}],
            [{"name": "Travis", "state_of_birth": "NY", "money": 0},{"name": "Mike", "state_of_birth": "NY", "money": 0}]
            ]

Answer 1

You can use itertools like this:

import itertools

def func(x):
    return tuple([x['state_of_birth'], x['money'] != 0])

desired_output = list(list(v) for _,v in itertools.groupby(sorted(users, key=func), func))

group_by function is a generator that produces key and value . The key is derived from the key_function that we're passing to the itertools.groupb_by() . In your case, having the keys in not important, that's why it is ignore in for _, v .

Output:

[{'name': 'John', 'state_of_birth': 'CA', 'money': 0}, {'name': 'Bill', 'state_of_birth': 'CA', 'money': 0}]
[{'name': 'Andrew', 'state_of_birth': 'CA', 'money': 300}]
[{'name': 'Travis', 'state_of_birth': 'NY', 'money': 0}, {'name': 'Mike', 'state_of_birth': 'NY', 'money': 0}]
[{'name': 'Scott', 'state_of_birth': 'OR', 'money': 20}]

Answer 2

code:

users = [
            {"name": "John", "state_of_birth": "CA", "money": 0},
            {"name": "Andrew", "state_of_birth": "CA", "money": 300},
            {"name": "Scott", "state_of_birth": "OR", "money": 20},
            {"name": "Travis", "state_of_birth": "NY", "money": 0},
            {"name": "Bill", "state_of_birth": "CA", "money": 0},
            {"name": "Mike", "state_of_birth": "NY", "money": 0}
        ]

result = {}
for user in users:
    key = (user["state_of_birth"],user["money"])
    if key in result:
        result[key].extend([user])
    else:
        result[key] = [user]
for _,v in result.items():
    print(v)

result:

[{'name': 'John', 'state_of_birth': 'CA', 'money': 0}, {'name': 'Bill', 'state_of_birth': 'CA', 'money': 0}]
[{'name': 'Andrew', 'state_of_birth': 'CA', 'money': 300}]
[{'name': 'Scott', 'state_of_birth': 'OR', 'money': 20}]
[{'name': 'Travis', 'state_of_birth': 'NY', 'money': 0}, {'name': 'Mike', 'state_of_birth': 'NY', 'money': 0}]

Answer 3

If I understand the question right, you have a structure that is List[Dict] and you want to get a List[List[Dict]] where the inner list contain dictionaries that have the same state_of_birth and money > 0 boolean.

I would say the easiest solution is actually to use pandas

import pandas as pd

users = [
            {"name": "John", "state_of_birth": "CA", "money": 0},
            {"name": "Andrew", "state_of_birth": "CA", "money": 300},
            {"name": "Scott", "state_of_birth": "OR", "money": 20},
            {"name": "Travis", "state_of_birth": "NY", "money": 0},
            {"name": "Bill", "state_of_birth": "CA", "money": 0},
            {"name": "Mike", "state_of_birth": "NY", "money": 0}
        ]

df = pd.DataFrame.from_records(users)

# we need a column to indicate if money > 0
df["money_bool"] = df["money"] > 0

# groupby gives you an iterator of Tuple[key, sub-dataframe]
# dfs now holds a list of your grouped dataframes
dfs = [tup[1] for tup in df.groupby(["state_of_birth", "money_bool"])]

# you can now drop the money_bool column if you want
dfs = [df.drop("money_bool", axis=1) for df in dfs]

desired_output = [df.to_dict("records") for df in dfs]

Depending on the context of the problem, you may be better off staying in dataframe/tabular format

Answer 4

You need to make sure that the input to the groupby function is sorted. You can use the same key function as for grouping:

users = [
            {"name": "John", "state_of_birth": "CA", "money": 0},
            {"name": "Andrew", "state_of_birth": "CA", "money": 300},
            {"name": "Scott", "state_of_birth": "OR", "money": 20},
            {"name": "Travis", "state_of_birth": "NY", "money": 0},
            {"name": "Bill", "state_of_birth": "CA", "money": 0},
            {"name": "Mike", "state_of_birth": "NY", "money": 0}
        ]

def selector(item): return (item.get('state_of_birth'), item.get('money') != 0)
sorted_users = sorted(users, key=selector)
result = [list(group) for _, group in groupby(sorted_users, selector) ]

Output:

[
    [{'name': 'John', 'state_of_birth': 'CA', 'money': 0}, {'name': 'Bill', 'state_of_birth': 'CA', 'money': 0}],
    [{'name': 'Andrew', 'state_of_birth': 'CA', 'money': 300}], 
    [{'name': 'Travis', 'state_of_birth': 'NY', 'money': 0}, {'name': 'Mike', 'state_of_birth': 'NY', 'money': 0}],
    [{'name': 'Scott', 'state_of_birth': 'OR', 'money': 20}]
]

Answer 5

Although its name seems like it should be the way to go, itertools.groupby is not the correct function to use because it requires the data to be pre-sorted. Sorting brings your time complexity to O(n log(n)) for an algorithm that should be O(n).

To put that in perspective, if you have a million records to sort, instead of a million iterations, you now have 20 million iterations if you use groupby instead of a loop and dict. That's a pretty significant performance penalty.

If groupby was cleaner to write or didn't have an import, it might be justifiable, but it's less readable than a simpler approach using a plain loop and dictionary.

Pandas is fine, but there's really no reason to use it unless you're already doing so. It's like bringing in a space shuttle to grill a zucchini.

You can use defaultdict and a loop:

from collections import defaultdict
from pprint import pprint

users = [
    {"name": "John", "state_of_birth": "CA", "money": 0},
    {"name": "Andrew", "state_of_birth": "CA", "money": 300},
    {"name": "Scott", "state_of_birth": "OR", "money": 20},
    {"name": "Travis", "state_of_birth": "NY", "money": 0},
    {"name": "Bill", "state_of_birth": "CA", "money": 0},
    {"name": "Mike", "state_of_birth": "NY", "money": 0},
]

grouped = defaultdict(list)
groupby = "state_of_birth", "money"

for user in users:
    grouped[tuple([user[k] for k in groupby])].append(user)

pprint([*grouped.values()])

If you want "money is nonzero" rather than just the "money" value itself, you can use a custom grouping function:

grouped = defaultdict(list)

def group_by(x):
    return x["state_of_birth"], x["money"] != 0

for user in users:
    grouped[group_by(user)].append(user)

result = [*grouped.values()]

or inline the logic:

grouped = defaultdict(list)

for user in users:
    grouped[user["state_of_birth"], user["money"] != 0].append(user)

result = [*grouped.values()]