Python pandas: conditionally select a uniform sample from a dataframe
<p>Say I have a dataframe as such</p>

<pre><code>category1  category2   other_col   another_col ....
a          1
a          2
a          2        
a          3
a          3
a          1
b          10
b          10
b          10
b          11
b          11
b          11
</code></pre>

<p>I want to obtain a sample from my dataframe so that <code>category1</code> a uniform number of times. I'm assuming that there are an equal number of each type in <code>category1</code>. I know that this can be done with pandas using <code>pandas.sample()</code>. However, I also want to ensure  that that sample I select has <code>category2</code> equally represented as well. So, for example, if I have a sample size of 5, I would want something such as:</p>

<pre><code>a  1
a  2
b  10
b  11
b  10
</code></pre>

<p>I would <em>not</em> want something such as:</p>

<pre><code>a 1
a 1
b 10
b 10
b 10
</code></pre>

<p>While this is a valid random sample of <code>n=4</code>, it would not meet my requirements as I want to vary as much as possible the types of <code>category2</code>.</p>

<p>Notice that in the first example, because <code>a</code> was only sampled twice, that <code>3</code> was not not represented from <code>category2</code>. This is okay. The goal is to just as uniformly as possible, represent that sample data.</p>

<p>If it helps to provide a clearer example, one could thing having the categories <code>fruit</code>, <code>vegetables</code>, <code>meat</code>, <code>grains</code>, <code>junk</code>. In a sample size of 10, I would want as much as possible to represent each category. So ideally, 2 of each. Then each of those 2 selected rows belonging to the chosen categories would have subcategories that are also represented as uniformly as possible. So, for example, fruit could have a subcategories of red_fruits, yellow_fruits, etc. For the 2 fruit categories that are selected of the 10, red_fruits and yellow_fruits would both represented in the sample. Of course, if we had larger sample size, we would include more of the subcategories of fruit (green_fruits, blue_fruits, etc.).</p>

<p>Trick is building up a balanced array.  I provided a clumsy way of doing it.  Then cycle through the groupby object sampling by referencing the balanced array.</p>

<pre><code>def rep_sample(df, col, n, *args, **kwargs):
    nu = df[col].nunique()
    m = len(df)
    mpb = n // nu
    mku = n - mpb * nu
    fills = np.zeros(nu)
    fills[:mku] = 1

    sample_sizes = (np.ones(nu) * mpb + fills).astype(int)

    gb = df.groupby(col)

    sample = lambda sub_df, i: sub_df.sample(sample_sizes[i], *args, **kwargs)

    subs = [sample(sub_df, i) for i, (_, sub_df) in enumerate(gb)]

    return pd.concat(subs)
</code></pre>

<hr>

<h3>Demonstration</h3>

<pre><code>rep_sample(df, 'category1', 5)
</code></pre>

<p><a href="http://i.stack.imgur.com/cCoqj.png" rel="nofollow"><img src="http://i.stack.imgur.com/cCoqj.png" alt="enter image description here"></a></p>

python pandas dataframe uniform