遍历字典以创建列表

Question

我在一个名为favoriteColors的MongoDB集合中有以下4个字典：

{ "name" : "Johnny", "color" : "green" }
{ "name" : "Steve", "color" : "blue" },
{ "name" : "Ben", "color" : "red" },
{ "name" : "Timmy", "color" : "cyan" }

我想创建一个有序的颜色值列表，该列表与另一个有序列表匹配。

例如，如果我有列表["Johnny", "Steve", "Ben", "Johnny"]，新列表将是["green", "blue", "red", "green"]。

如果我有列表["Steve", "Steve", "Ben", "Ben", "Johnny"]，新列表将是["blue", "blue", "red", "red", "green"]。

使用Python和/或PyMongo，有什么好的方法可以完成这个任务。这是我目前的代码，但它没有识别重复项。

name_list = ["Steve", "Steve", "Ben", "Ben", "Johnny"]

color_list = []
for document in db.favoriteColors.aggregate([
    {"$match": {"name": {"$in": name_list }}},
    {"$project": {"color": 1}}
]):
    for k, v in document.iteritems():
        color_list.append(v)

print color_list
# ["blue", "red", "green"]

- Johnny Metz

---

你能否创建一个 'name' -> 'color' 的映射，然后使用该映射与 name_list 一起构建 color_list？我对 MongoDB 不够了解，不知道是否可以高效地完成，但似乎是可以的... - mgilson

Answer 1

实际上，我们可以使用聚合框架和客户端处理来高效地完成这个任务。

import pymongo


client = pymongo.MongoClient()
db = client.test # Or whatever is your database
favoriteColors = db.favoriteColors
first_list = ['Johnny', 'Steve', 'Ben', 'Johnny']

cursor = favoriteColors.aggregate([
    {'$match': {'name': {'$in': first_list}}}, 
    {'$project': {'part': {'$map': {
        'input': first_list, 
        'as': 'inp', 
        'in': {
            '$cond': [
                {'$eq': [ '$$inp', '$name']}, 
                '$color', 
                None
            ]
        }
    }}}},
    {'$group': {'_id': None, 'data': {'$push': '$part'}}}
])

由于我们按照 None 进行 $group，所以我们的光标包含一个文档，我们可以使用 next 检索该文档。实际上，我们可以使用 print(list(cursor)) 来验证这一点。

>>> import pprint
>>> pprint.pprint(list(cursor))
[{'_id': None,
  'data': [['green', None, None, 'green'],
           [None, 'blue', None, None],
           [None, None, 'red', None]]}]

从这里开始，我们需要使用zip来解包文档中的“data”字段，使用chain.from_iterable链接输入，并过滤掉None元素。

from itertools import chain

result = [item 
          for item in chain.from_iterable(zip(*next(cursor)['data']))
          if item is not None]

这将返回：

>>> result
['green', 'blue', 'red', 'green']

- styvane

---

你能分享一下你的答案在足够多的查询下的基准测试结果吗？看起来很高效。谢谢。 - Chiheb Nexus

这个看起来好像行得通，但是看起来太复杂了。肯定有一种更简单的方法可以使用聚合框架来做这个。为什么需要在"$project"下面有那么多的条目呢？ - Johnny Metz

@jcmetz21 我认为这并不复杂，如果输出列表中的元素不需要按照这个确切的顺序出现，我会添加解释和可能的另一种方法，但这并不会更加简单。 - styvane

元素需要按照精确顺序排列，这样才能看起来好。谢谢！ - Johnny Metz

@jcmetz21 我的答案不仅看起来好，而且是完成这个任务的最佳方式。在MongoDB 3.4中可能还有其他方法可以实现。 - styvane

Answer 2

如果数据集很小，您可以将字典合并为一个新的字典。

在Python3中，您可以像这样做：

names = ["Steve", "Steve", "Ben", "Ben", "Johnny"]
favorites = {d["name"]: d["color"] for d in db.favoriteColors.find()}
colors = [favorites[name] for name in names]
print(colors)

已更新

正如Styvane所提到的，我忘记在Collection上调用find方法。答案已相应更新。

- Hans Then

---

集合对象不可迭代。这将会因为一个美丽的TypeError而悲惨地失败。 - styvane

如果字典是在Python中且要迭代的数据很少，那么这是一种干净的方法。 - Johnny Metz

@jcmetz21，我没看出这是“一种干净的方式”。db.favoriteColors是一个Collection对象，不实现迭代器协议，因此这个查询将返回一个TypeError，正如我所提到的。 - styvane

如果我将字典放入列表中并遍历该列表，则可以工作。但是如果我们从MongoDB抓取字典，我确实同意这是可怕的做法。 - Johnny Metz

Answer 3

您还可以从当前字典中创建一个新的dict，其中每个dict["name"]的值将与dict["color"]的值相关联。

例如：新的字典将如下所示：

{"Jhonny": "green", "Steve": "blue"}

你可以使用像下面这个例子一样的函数，接受许多参数并返回所需的列表（如果输入列表中有任何默认字典中不存在的名称，则还会附加None）：

这是我的例子：

a = { "name" : "Johnny", "color" : "green" }
b = { "name" : "Steve", "color" : "blue" }
c = { "name" : "Ben", "color" : "red" }
d = { "name" : "Timmy", "color" : "cyan" }

my_list = ["Steve", "Steve", "Ben", "Ben", "Johnny"]

def iter_func(my_list = list, *args):
    ne = {k["name"]:k["color"] for k in args}
    return [ne[k] if k in ne.keys() else None for k in my_list]

输出：

print(iter_func(my_list, a,b,c,d))
>>> ['blue', 'blue', 'red', 'red', 'green']

带有None值的示例：

a = { "name" : "Johnny", "color" : "green" }
b = { "name" : "Steve", "color" : "blue" }
c = { "name" : "Ben", "color" : "red" }
d = { "name" : "Timmy", "color" : "cyan" }

my_list = ["Steve", "Steve", "Alex", "Ben", "Ben", "Johnny", "Mark"]

def iter_func(my_list = list, *args):
    ne = {k["name"]:k["color"] for k in args}
    return [ne[k] if k in ne.keys() else None for k in my_list]

输出：

print(iter_func(my_list, a,b,c,d))
>>> ['blue', 'blue', None, 'red', 'red', 'green', None]

- Chiheb Nexus

---

这不是正确的做法。如果有一百万个文档符合您的查询，那该怎么办？ - styvane

OP没有说他的查询像你所说的那样巨大。但是感谢您的评论。我会尝试改进我的答案。 - Chiheb Nexus