根据熊猫数据框的索引定义列表

Question

I have a pandas dataframe, and one of the columns has date values as strings (like "2014-01-01"). I would like to define a different list for each year that is present in the column, where the elements of the list are the index of the row in which the year is found in the dataframe.

Here's what I've tried:

import pandas as pd    

df = pd.DataFrame(["2014-01-01","2013-01-01","2014-02-02", "2012-08-09"])
df = df.values.flatten().tolist()

for i in range(len(df)):
    df[i] = df[i][0:4]

y2012 = []; y2013 = []; y2014 = []

for i in range(len(df)):
    if df[i] == "2012":
        y2012.append(i)
    elif df[i] == "2013":
        y2013.append(i)
    else:
        y2014.append(i)

print y2014 # [0, 2]
print y2013 # [1]
print y2012 # [3]

Does anyone know a better way of doing this? This way works fine, but I have a lot of years, so I have to manually define each variable and then run it through the for loop, and so the code gets really long. I was trying to use groupby in pandas, but I couldn't seem to get it to work.

Thank you so much for any help!

Answer 1

Scan through the original DataFrame values and parse out the year. Given, that, add the index into a defaultdict. That is, the following code creates a dict , one item per year. The value for a specific year is a list of the rows in which the year is found in the dataframe.

A defaultdict sounds scary, but it's just a dictionary. In this case, each value is a list. If we append to a nonexistent value, then it gets spontaneously created. Convenient!

source

from collections import defaultdict
import pandas as pd    

df = pd.DataFrame(["2014-01-01","2013-01-01","2014-02-02", "2012-08-09"])
# df = df.values.flatten().tolist()

dindex = defaultdict(list)
for index,dateval in enumerate(df.values):
    year = dateval[0].split('-')[0]
    dindex[year].append(index)

assert dindex == {'2014': [0, 2], '2013': [1], '2012': [3]}
print dindex

output

defaultdict(<type 'list'>, {'2014': [0, 2], '2013': [1], '2012': [3]})

Answer 2

Pandas is awesome for this kind of thing, so don't be so hasty to turn your dataframe back into lists right away.

The trick here lies in the .apply() method and the .groupby() method.

1. Take a dataframe that has strings with ISO formatted dates in it
2. parse the column containing the date strings into datetime objects
3. Create another column of years using the datetime.year attribute of the items in the datetime column
4. Group the dataframe by the new year column
5. Iterate over the groupby object and extract your column

Here's some code for you to play with and grok:

import pandas
import dateutil

df = pd.DataFrame({'strings': ["2014-01-01","2013-01-01","2014-02-02", "2012-08-09"]})
df['datetimes'] = df['strings'].apply(dateutil.parser.parse)
df['year'] = df['datetimes'].apply(lambda x: x.year)
grouped_data= df.groupby('year')

lists_by_year = {}
for year, data in grouped_data
    lists_by_year [year] = list(data['strings'])

Which gives us a dictionary of lists, where the key is the year and the contents is a list of strings with that year.

print lists_by_year 

{2012: ['2012-08-09'],
 2013: ['2013-01-01'],
 2014: ['2014-01-01', '2014-02-02']}

Answer 3

As it turns out

df.groupby('A') #is just syntactical sugar for df.groupby(df['A'])

This means that all you have to do to group by year is leverage the apply function and re-work the syntax

Solution

getYear = lambda x:x.split("-")[0]
yearGroups = df.groupby(df["dates"].apply(getYear))

Output

for key,group in yearGroups: 
    print key

2012
2013
2014