python - 累计和与随机抽奖号码比较后返回列名

Question

我有一个 DataFrame，其中列的总和为 1，如下所示:

Out[]: 
cod_mun  ws_1    1_3      4      5_7       8      9_10    11     12    13     14     15     nd
1100015 0.1379  0.273   0.2199  0.1816  0.0566  0.0447  0.0617  0.0015  0   0.0021  0.0074  0.0137
1100023 0.1132  0.2009  0.185   0.2161  0.1036  0.0521  0.0885  0.0044  0.0038  0.0061  0.0181  0.0082

我想抽一个随机数

import random
prob = random.random()

然后我想将这个数字与从左到右列的累积总和进行比较，然后返回列的标题。

例如，如果 prob = 0.24 阈值将在第二列中达到 0.27，0.1379 + 0.273 > 0.24那么我需要返回列的名称.

是否可以在不使用 15 个 elif 的情况下做到这一点？

这样:

if prob < df.iloc[0]['ws_1']:
    return 'ws_1'
elif prob < df.iloc[0]['ws_1'] + df.iloc[0]['1_3']
    return '1_3'
elif ...

Answer 1

我想你可以数DataFrame.cumsum , 与 prob 比较并通过 idxmax 获得具有 True 值的第一列:

df.set_index('cod_mun', inplace=True)

prob = 0.24 

print (df.cumsum(axis=1))
           ws_1     1_3       4     5_7       8    9_10      11      12  \
cod_mun                                                                   
1100015  0.1379  0.4109  0.6308  0.8124  0.8690  0.9137  0.9754  0.9769   
1100023  0.1132  0.3141  0.4991  0.7152  0.8188  0.8709  0.9594  0.9638   

             13      14      15      nd  
cod_mun                                  
1100015  0.9769  0.9790  0.9864  1.0001  
1100023  0.9676  0.9737  0.9918  1.0000  

print (df.cumsum(axis=1) > prob)
          ws_1   1_3     4   5_7     8  9_10    11    12    13    14    15  \
cod_mun                                                                      
1100015  False  True  True  True  True  True  True  True  True  True  True   
1100023  False  True  True  True  True  True  True  True  True  True  True   

           nd  
cod_mun        
1100015  True  
1100023  True

print ((df.cumsum(axis=1) > prob).idxmax(axis=1))
cod_mun
1100015    1_3
1100023    1_3
dtype: object