将累积和与随机抽取的数字进行比较后返回列的名称

时间:2016-09-12 12:38:06

标签: python pandas random

我有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()

然后我想将这个数字与列从左到右的累积总和进行比较,然后返回列'heading

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

没有使用15 elif s?

可以做到这一点

这样:

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 ...

1 个答案:

答案 0 :(得分:4)

我认为您可以计算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