随机获取加权值的Python字典键

35 投票
10 回答
18763 浏览
提问于 2025-04-15 12:33

我有一个字典,每个键对应一个长度不一的列表,比如:

d = {
 'a': [1, 3, 2],
 'b': [6],
 'c': [0, 0]
}

有没有什么简单的方法可以根据值的长度来随机获取字典中的一个键?使用random.choice(d.keys())会让每个键的被选中几率一样,但在上面的例子中,我希望'a'大约能被选中一半的时间。

10 个回答

9

不需要创建一个新的、可能很大的包含重复值的列表:

def select_weighted(d):
   offset = random.randint(0, sum(d.itervalues())-1)
   for k, v in d.iteritems():
      if offset < v:
         return k
      offset -= v
17

你是否总是知道字典中有多少个值?如果是这样的话,下面这个算法可能会很简单,适合在你想从一个有序列表中随机选择一些项目时使用:

  1. 遍历你的键列表。
  2. 生成一个在0到1之间均匀分布的随机值(就像“掷骰子”一样)。
  3. 假设这个键有N_VALS个值与之相关,并且整个字典中总共有TOTAL_VALS个值,那么以N_VALS / N_REMAINING的概率接受这个键,其中N_REMAINING是列表中剩下的项目数量。

这个算法的好处是不用生成新的列表,这在字典很大的时候特别重要。你的程序只需要遍历K个键来计算总数,再遍历一次键,平均来说会到达一半的地方,还有生成0到1之间随机数的成本。生成这样的随机数在编程中是非常常见的,所以大多数编程语言都有快速实现这个功能的方法。在Python中,随机数生成器使用的是梅森旋转算法的C语言实现,速度应该非常快。此外,文档还声称这个实现是线程安全的。

下面是代码。如果你想使用更Pythonic的特性,我相信你可以把它整理得更好:

#!/usr/bin/python

import random

def select_weighted( d ):
   # calculate total
   total = 0
   for key in d:
      total = total + len(d[key])
   accept_prob = float( 1.0 / total )

   # pick a weighted value from d
   n_seen = 0
   for key in d:
      current_key = key
      for val in d[key]:
         dice_roll = random.random()
         accept_prob = float( 1.0 / ( total - n_seen ) )
         n_seen = n_seen + 1
         if dice_roll <= accept_prob:
            return current_key

dict = {
   'a': [1, 3, 2],
   'b': [6],
   'c': [0, 0]
}

counts = {}
for key in dict:
   counts[key] = 0

for s in range(1,100000):
   k = select_weighted(dict)
   counts[k] = counts[k] + 1

print counts

运行这个100次后,我得到的选择键的次数是:

{'a': 49801, 'c': 33548, 'b': 16650}

这些结果和你预期的值相当接近:

{'a': 0.5, 'c': 0.33333333333333331, 'b': 0.16666666666666666}

编辑:Miles指出了我原始实现中的一个严重错误,已经修正了。对此我感到抱歉!

35

这样做是可以的:

random.choice([k for k in d for x in d[k]])

撰写回答