Writing variable sized arrays to Pandas cells
<p>I have a large data set and I want to do a convolution calculation using multiple rows that match a criteria. I need to calculate a vector for each row first, and I thought it would be more efficient to store my vector in a dataframe column so I could try and avoid a for loop when I do the convolution. Trouble is, the vectors are variable length and I can't figure out how to do it.</p>

<p>Here's a summary of my data:</p>

<pre><code>Date        State  Alloc P
2012-01-01  AK     3     0.5
2012-01-01  AL     4     0.3
â¦
</code></pre>

<p>Each state has a different Alloc and P value. Thereâs a row for every date and state and my dataframe is over 15,000 rows long.</p>

<p>For each entry, I want a vector that looks like this:</p>

<pre><code>[P, np.zeros(Alloc), 1-P]
</code></pre>

<p>I can't figure out how to set a new column like this. I've tried statements like:</p>

<pre><code>df['Test'] = [df['P'], np.zeros(df['Alloc'), 1 â df['P']]
</code></pre>

<p>but they don't work.</p>

<p>Does anyone have any ideas?</p>

<p>Thanks âº</p>

<p>Try:</p>

<pre><code>def get_vec(x):
    return [x.P] + np.zeros(x['Alloc']).tolist() + [1 - x.P]

df.apply(get_vec, axis=1)

0         [0.5, 0.0, 0.0, 0.0, 0.5]
1    [0.3, 0.0, 0.0, 0.0, 0.0, 0.7]
dtype: object
</code></pre>

<hr>

<pre><code>df['Test'] = df.apply(get_vec, axis=1)
df
</code></pre>

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

python pandas