Return name of column after comparing cummulative sum with random drawn number
<p>I have a <code>DataFrame</code> in which the sum of the columns is <code>1</code>, like so:</p>

<pre><code>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
</code></pre>

<p>I want to draw a random number </p>

<pre><code>import random
prob = random.random()
</code></pre>

<p>And then I want to compare such number with the cummulative sum of the columns <strong>from left to right</strong> and then return the columns' <code>heading</code>.</p>

<p>For example, if <code>prob = 0.24</code> the threshold would reach 0.27 in the second column, <code>0.1379 + 0.273 &gt; 0.24</code>Then I would need to return the name of the column.</p>

<p>It it possible to do that WITHOUT a using 15 <code>elif</code>s?</p>

<p>Such that:</p>

<pre><code>if prob &lt; df.iloc[0]['ws_1']:
    return 'ws_1'
elif prob &lt; df.iloc[0]['ws_1'] + df.iloc[0]['1_3']
    return '1_3'
elif ...
</code></pre>

<p>I think you can count <a href="http://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.cumsum.html" rel="nofollow"><code>DataFrame.cumsum</code></a>, compare with <code>prob</code> and get first column with <code>True</code> value by <a href="http://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.idxmax.html" rel="nofollow"><code>idxmax</code></a>:</p>

<pre><code>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) &gt; 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) &gt; prob).idxmax(axis=1))
cod_mun
1100015    1_3
1100023    1_3
dtype: object
</code></pre>

python pandas random