Chen Chen - 6 months ago 11
Python Question

Manipulate column data to new format by using Python in Pandas

I am trying to use pandas to manipulate my data. I have data that looks like this:

CarModel ProductionData ProductionYear
BMWX1 55000 2005
Accord 100000 2005
BMWX1 34000 2006
Accord 110000 2006
BMWX1 43000 2007
Accord 105000 2007


How can I make the data look like this?


CarModel NewColumn
BMWX1 [2005.0, 2006.0, 55000]; [2006.0, 2007.0, 34000]; [2007.0, 2008.0, 43000]
Accord [2005.0, 2006.0, 100000]; [2006.0, 2007.0, 110000]; [2007.0, 2008.0, 105000]

Answer

Setup

text = """CarModel ProductionData ProductionYear
BMWX1    55000          2005
Accord   100000         2005
BMWX1    34000          2006
Accord   110000         2006
BMWX1    43000          2007
Accord   105000         2007"""

df = pd.read_csv(StringIO(text), delim_whitespace=1)

Solution

gb = df.set_index('CarModel').groupby(level=0)

def proc_df(df):
    # Add this column becuase OP has it in final output
    df['Year2'] = df.ProductionYear + 1

    columns = ['ProductionYear', 'Year2', 'ProductionData']

    # Return ndarray gets flattened to string when returned via apply
    return df[columns].values

gb.apply(proc_df)

Looks like:

CarModel
Accord    [[2005, 2006, 100000], [2006, 2007, 110000], [...
BMWX1     [[2005, 2006, 55000], [2006, 2007, 34000], [20...
dtype: object