То, что вы можете сделать, - это использовать np.random.choice с такими же вероятностями, как у вас, но для полного размера ваших данных. Затем reindex df с новым заказом, полученным от np.random.choice. Используйте cumsum для веса столбца и, наконец, возвращайте только индекс, пока он не достигнет желаемого значения.
def weighted_budgeted_random_sample_all(df, budget):
random_index_order = np.random.choice( df.index, size = len(df),
p = df.prob, replace = False)
s = df.reindex(random_index_order).weight.cumsum()
return s[s <= budget].index.values
Теперь проблема этого метода заключается в том, что при df, как в вопросе, и budget из 10, некоторые решения имеют только индекс 1 или 2, потому что если random_index_order равно [2,1,0] или [1,2,0], тогда cumsum выше 10 во втором ряду.
См. С Counter, использование tuple и np.sort просто для того, чтобы Counter заработало и чтобы было легче увидеть результат:
from collections import Counter
print (Counter([ tuple(np.sort(weighted_budgeted_random_sample_all(df,10)))
for i in range(1000)]))
# Counter({(0, 1): 167, (0, 2): 111, (1,): 390, (2,): 332})
Как вы можете видеть, некоторые розыгрыши были в порядке с 2 и 3 в качестве первых 2 значений, и результат равен только 2 или 3, потому что сумма их весов равна 11.
Но на самом деле, если вы попробуете то же самое с бюджетом 11, вы получите ожидаемый результат:
print (Counter([ tuple(np.sort(weighted_budgeted_random_sample_all(df,11)))
for i in range(1000)]))
# Counter({(0, 1): 169, (0, 2): 111, (1, 2): 720})
Здесь вы найдете три возможных набора и тот факт, что набор {1,2} получается чаще, имеет смысл.
Я видел, что вы пересмотрели свой вопрос после комментария, в котором говорилось, что вы будете работать над подходом «один элемент за один раз» . Я верю, что это повлияет на общие вероятности, но я не знаю достаточно вероятностей, чтобы сказать, почему. Если вы действительно хотите, то я предполагаю, что вы можете объединить свой подход и мой, чтобы выиграть время:
def weighted_budgeted_random_sample_mixed(df, budget):
ids = []
total = 0
dftemp = df.copy()
while total < budget:
remaining = budget - total
dftemp = dftemp[dftemp.weight <= remaining]
# Stop if there are no records with small enough weight.
if dftemp.shape[0] == 0:
break
# New order
new_index = np.random.choice( dftemp.index, size = len(dftemp),
p = (dftemp.prob/dftemp.prob.sum()),
replace = False)
s = dftemp.reindex(new_index).weight.cumsum()
#select only the necessary rows
s = s[s <= remaining]
total += s.max() #last value in s which is less than remaining
dftemp.drop(s.index, inplace=True)
ids += s.index.tolist()
return ids
Теперь для сравнения с вашим методом с точки зрения результата:
#your approach
print (Counter([ tuple(np.sort(weighted_budgeted_random_sample(df,10)))
for i in range(1000)]))
#Counter({(0, 1): 546, (0, 2): 454})
#mixed approach
print (Counter([ tuple(np.sort(weighted_budgeted_random_sample_mixed(df,10)))
for i in range(1000)])
#Counter({(0, 1): 554, (0, 2): 446})
Как вы можете видеть, результат довольно схожий, и смешанный подход должен быть быстрее на больших фреймах данных, поскольку он минимизирует зацикливание в while