DataFrame conditional update of multiple columns

Question

I'm looking for a more elegant and Pythonic code for the following problem:

In column start I have a release week for every item (from 1 to 4), I add W1,..., W4 columns with ones.

Next I want to update columns in this way (basically, replace the ones in the release week and the week before and week after with zeros):

   start  W1  W2  W3  W4
      1   0   0   1   1
      2   0   0   0   1
      3   1   0   0   0
      4   1   1   0   0

I'm doing it right now with this:

import pandas as pd

data = {'start': [1,2,3,4]}
df = pd.DataFrame(data)
for i in range(1,4+1):
    df['W'+str(i)] = 1

for index, i in enumerate(df['start']):
    if i==1:
        df.ix[index, 'W1'] = 0
        df.ix[index, 'W2'] = 0
    elif i==4:
        df.ix[index, 'W3'] = 0
        df.ix[index, 'W4'] = 0
    else:
        df.ix[index, 'W'+str(i-1)] = 0
        df.ix[index, 'W'+str(i)] = 0
        df.ix[index, 'W'+str(i+1)] = 0

Thanks for your help. Solution with map is slower but for me more readable (personal taste) and shorter. Also it don't use numpy. One thing: when I added more rows, it returns error. Probably, there is an easy fix, but I didn't look for. Once more than for your help and for sharing different approach. — Quant Christo
– Quant Christo, Commented Jan 23, 2018 at 9:02
You can upvote all solutions you found helpful. By the way, you should compare the performances of the solution you accepted v/s mine... you'll find there is a staggering difference. — cs95
– cs95, Commented Dec 28, 2018 at 5:14

joaquin · Accepted Answer · 2018-01-21 19:38:34Z

1

From your df

df = df.astype(int)

start   W1  W2  W3  W4
0   1   1   1   1   1
1   2   1   1   1   1
2   3   1   1   1   1
3   4   1   1   1   1

You can apply map to df to recalculate values using the function:

def func(pivot):
    return [1 if abs(col-pivot) > 1 else 0 for col in [1,2,3,4]]

This was my first option for map, unnecessarily complex:

df['W1'], df['W2'], df['W3'], df['W4'] = zip(*df['start'].map(func))

This one, from @QuantChristo, is much better

weeks = ['W1','W2','W3','W4']
df[weeks] = df['start'].map(func)

In both cases you get df:

start   W1  W2  W3  W4
0   1   0   0   1   1
1   2   0   0   0   1
2   3   1   0   0   0
3   4   1   1   0   0

edited Jan 21, 2018 at 19:38

answered Jan 21, 2018 at 14:10

joaquin

86k31 gold badges146 silver badges155 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Quant Christo Over a year ago

Thank you for your help! I've slightly changed code to less intimitading to me (no zip and *)

def func(pivot):      return [1 if abs(col-pivot) > 1 else 0 for col in range(1,4+1)]  weeks = ['W1','W2','W3','W4'] df[weeks] = df['start'].map(func)

joaquin Over a year ago

In fact, looks simpler and better. I updated the answer to take your code into account

Quant Christo Over a year ago

With lambda function, it could even be one-liner (or two-liner if you count weeks definition). One advantage of weeks, is that you can easily add more weeks (e.g. 52) with list comprehension.

cs95 · Accepted Answer · 2018-01-22 12:14:50Z

Perform broadcasted numpy comparison to obtain a mask, and just set the values at the corresponding indices to 0.

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

i = df.index.values
j = np.arange(1, len(df) + 1)[:, None]

df.values[(i - 1 <= j) & (j <= i + 1)] = 0
df

       W1  W2  W3  W4
start                
1       0   0   1   1
2       0   0   0   1
3       1   0   0   0
4       1   1   0   0

Details

i
array([1, 2, 3, 4]) 

j 
array([[1],
       [2],
       [3],
       [4]])

First, compute the mask -

m = (i - 1 <= j) & (j <= i + 1)
m

array([[ True,  True, False, False],
       [ True,  True,  True, False],
       [False,  True,  True,  True],
       [False, False,  True,  True]], dtype=bool)

The mask m is for the entire dataframe. Just index on values and set the cells to 0 -

df.values[m] = 0

To reset the index, use reset_index -

df.reset_index()

   start  W1  W2  W3  W4
0      1   0   0   1   1
1      2   0   0   0   1
2      3   1   0   0   0
3      4   1   1   0   0

Collectives™ on Stack Overflow

DataFrame conditional update of multiple columns

2 Answers 2

3 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

3 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related