How to increment index per column / group

Viewed 235

I need to format the dataframe from :

| country   | county   | city   | street    |
|-----------|----------|--------|-----------|
| country 1 | county 1 | city 1 | street 1  |
| country 1 | county 1 | city 1 | street 2  |
| country 1 | county 1 | city 2 | street 3  |
| country 2 | county 2 | city 3 | street 4  |
| country 2 | county 2 | city 3 | street 5  |
| country 3 | county 3 | city 4 | street 6  |
| country 3 | county 4 | city 5 | street 7  |
| country 3 | county 4 | city 6 | street 8  |
| country 3 | county 4 | city 6 | street 9  |
| country 3 | county 4 | city 6 | street 10 |

to

| country   | county   | city   | street    | count |
|-----------|----------|--------|-----------|-------|
| country 1 |          |        |           | 3     |
|           | county 1 |        |           | 3     |
|           |          | city 1 |           | 2     |
|           |          |        | street 1  | 1     |
|           |          |        | street 2  | 1     |
|           |          | city 2 |           | 1     |
|           |          |        | street 3  | 1     |
| country 2 |          |        |           | 2     |
|           | county 2 |        |           | 2     |
|           |          | city 3 |           | 2     |
|           |          |        | street 4  | 1     |
|           |          |        | street 5  | 1     |
| country 3 |          |        |           | 5     |
|           | county 3 |        |           | 1     |
|           |          | city 4 |           | 1     |
|           |          |        | street 6  | 1     |
|           | county 4 |        |           | 4     |
|           |          | city 5 |           | 1     |
|           |          |        | street 7  | 1     |
|           |          | city 6 |           | 3     |
|           |          |        | street 8  | 1     |
|           |          |        | street 9  | 1     |
|           |          |        | street 10 | 1     |

The number of columns may vary.

I am managing the count with multiple groupby and trying the formatting in python without success. There is a way to do it only with pandas?

4 Answers

The solution I propose works fine for any number of columns but it's not a perfect solution for very big dataframes:

def add_column(col_name,col_value,count,tempdf):
    return tempdf.append({col_name:col_value,'count':count},ignore_index=True)

newdf = pd.DataFrame()
col_arr = df.columns.to_list()
col_dict ={}
for i,row in df.iterrows():
    for col in row.to_dict().keys():
        if current[col] != row[col]:
            for c in col_arr[col_arr.index(col):]:
                current[c] = row[c]
                newdf = add_column(c,row[c],df[lambda x: x[c]==row[c]].shape[0],newdf)

The result will be newdf:

      country    county    city     street  count
0   country 1       NaN     NaN        NaN    3.0
1         NaN  county 1     NaN        NaN    3.0
2         NaN       NaN  city 1        NaN    2.0
3         NaN       NaN     NaN   street 1    1.0
4         NaN       NaN     NaN   street 2    1.0
5         NaN       NaN  city 2        NaN    1.0
6         NaN       NaN     NaN   street 3    1.0
7   country 2       NaN     NaN        NaN    2.0
8         NaN  county 2     NaN        NaN    2.0
9         NaN       NaN  city 3        NaN    2.0
10        NaN       NaN     NaN   street 4    1.0
11        NaN       NaN     NaN   street 5    1.0
12  country 3       NaN     NaN        NaN    5.0
13        NaN  county 3     NaN        NaN    1.0
14        NaN       NaN  city 4        NaN    1.0
15        NaN       NaN     NaN   street 6    1.0
16        NaN  county 4     NaN        NaN    4.0
17        NaN       NaN  city 5        NaN    1.0
18        NaN       NaN     NaN   street 7    1.0
19        NaN       NaN  city 6        NaN    3.0
20        NaN       NaN     NaN   street 8    1.0
21        NaN       NaN     NaN   street 9    1.0
22        NaN       NaN     NaN  street 10    1.0

You can iterate over the columns themselves and rely on DataFrame.value_counts() to get you counts at the varying levels of nestedness. You'll need to play around with the index while doing this to get everything to realign correctly later, but in the end you just use pd.concat to stick these chunks together:

chunk_counts = []

for col in test_df.columns:
    counts = test_df.loc[:, :col].value_counts()
    n_empty_levels = test_df.columns.size - test_df.columns.get_loc(col) - 1
    empty_levels = [[""]] * n_empty_levels
    
    new_levels = [*counts.index.levels, *empty_levels]
    new_index = pd.MultiIndex.from_product(new_levels, names=test_df.columns)
    
    chunk_counts.append(counts.reindex(new_index))
    
final_series = (pd.concat(chunk_counts)
                .sort_index()
                .dropna()
                .astype(int)
                .rename("count"))

The repr appears fine if you print(final_series), but the multiindex doesnt have empty entries below each nested level (just appears that way with how MultiIndex displays. This becomes evident when we use reset_index. To put our series back into a frame requires maintain the format OP requested, we need to make a few more adjustments.

index_cols = final_series.index.names
final_df = final_series.reset_index()
final_df[index_cols] = final_df[index_cols].where(~final_df[index_cols].apply(pd.Series.duplicated))
final_df = final_df.fillna("")

print(final_df)

      Country    County    City     Street  count
0   Country 1                                   3
1              County 1                         3
2                        City 1                 2
3                                 Street 1      1
4                                 Street 2      1
5                        City 2                 1
6                                 Street 3      1
7   Country 2                                   2
8              County 2                         2
9                        City 3                 2
10                                Street 4      1
11                                Street 5      1
12  Country 3                                   5
13             County 3                         1
14                       City 4                 1
15                                Street 6      1
16             County 4                         4
17                       City 5                 1
18                                Street 7      1
19                       City 6                 3
20                               Street 10      1
21                                Street 8      1
22                                Street 9      1

You can first iterate over the rows of the original dataframe, and create a list of offsets. Then, with this list of offsets, you can create a new dataframe with the desired format.

Here is you dataframe

countries = ['Country %i' % ind for ind in [1,1,1,2,2,3,3,3,3,3]]
counties = ['County %i' % ind for ind in [1,1,1,2,2,3,4,4,4,4]]
cites = ['City %i' % ind for ind in [1,1,2,3,3,4,5,6,6,6]]
streets = ['Street %i' % ind for ind in [1,2,3,4,5,6,7,8,9,10]]

test_df = pd.DataFrame({'Country':countries,'County':counties,'City':cites,'Street':streets})

To create the list of offsets, you need

col_names = test_df.columns.tolist()

col_inds = [[] for _ in range(len(col_names))]
col_offsets = np.arange(len(col_names))
row_buff = None

# iterate over rows of the dataframe
# the array col_offsets of a row records the indices of the row entries in the output dataframe 
for ind,row in test_df.iterrows():

    # if it is the first row, the col_offsets is clearly [0,1,2,3]
    if ind == 0:
        row_buff = row.values
        for ind_offset,offset in enumerate(col_offsets):
            col_inds[ind_offset].append(offset)
    else:
        # if it is not the first row, 
        # find from which column the current row is different from the last row
        diff_ind = np.argmax(~(row.values==row_buff))
        last_offset = col_offsets.copy()
        for disp_count,disp_ind in enumerate(range(diff_ind,len(col_names))):
            col_offsets[disp_ind] = last_offset[-1] + disp_count + 1
            col_inds[disp_ind].append(col_offsets[disp_ind])
            
        row_buff = row.values

The list of offsets is col_inds, in this case, it is

[[0, 7, 12],
 [1, 8, 13, 16],
 [2, 5, 9, 14, 17, 19],
 [3, 4, 6, 10, 11, 15, 18, 20, 21, 22]]

Take the first sublist for example, the row indices of Country 1, Country 2, and Country 3 should be 0, 7, and 12, in the first column of the output dataframe.

With this col_inds, you can create the dataframe you want, with

max_ind = max([inner for item in col_inds for inner in item])
# initialize an empty string array 
placeholder = np.empty((max_ind+1,len(col_names)),dtype='S10')
placeholder[:,:] = ''

for ind,item in enumerate(col_inds):
    
    if ind < len(col_inds) - 1:
        df_col = test_df.groupby(col_names[:ind+1],as_index=False).agg({col_names[ind+1]:'count'}).iloc[:,ind]
    else:
        df_col = test_df.iloc[:,ind]
        
    for val_ind,val_pos in enumerate(item):
        placeholder[val_pos,ind] = df_col[val_ind]

# create new dataframe
new_df = pd.DataFrame({name:val for name,val in zip(col_names,placeholder.T)})
for col in col_names:
    new_df[col] = new_df[col].str.decode('utf-8') 

Here is another version and a variation. The first version is a bit simpler, but prints the result not quite as you specified it, but it should be fast, since it only requires four groupby statements a concat and a sort_index. The second version is a modification of the first version, which should still be fast, but the ouput looks like in your question:

cols= ['country', 'county', 'city', 'street']
dfs_result= list()
group_key= list()
cols_left= cols
while cols_left:
    count_key, *cols_left= cols_left
    group_key.append(count_key)
    df_counts= df.groupby(group_key).agg(count=(count_key, 'count'))
    df_counts.reset_index(drop=False, inplace=True)
    for col in cols:
        if col not in df_counts.columns:
            # create a fake column with NAs, so we have the same structure for
            # all count dfs
            df_counts[col]= np.NaN
    df_counts.set_index(cols, inplace=True)
    dfs_result.append(df_counts)

df_result= pd.concat(dfs_result, axis='index')
df_result.sort_index(inplace=True)
df_result

The result looks like this:

      country    county    city     street  count
0   country 1  county 1  city 1   street 1      1
1   country 1  county 1  city 1   street 2      1
2   country 1  county 1  city 1        NaN      2
3   country 1  county 1  city 2   street 3      1
4   country 1  county 1  city 2        NaN      1
5   country 1  county 1     NaN        NaN      3
6   country 1       NaN     NaN        NaN      3
7   country 2  county 2  city 3   street 4      1
8   country 2  county 2  city 3   street 5      1
9   country 2  county 2  city 3        NaN      2
10  country 2  county 2     NaN        NaN      2
11  country 2       NaN     NaN        NaN      2
12  country 3  county 3  city 4   street 6      1
13  country 3  county 3  city 4        NaN      1
14  country 3  county 3     NaN        NaN      1
15  country 3  county 4  city 5   street 7      1
16  country 3  county 4  city 5        NaN      1
17  country 3  county 4  city 6  street 10      1
18  country 3  county 4  city 6   street 8      1
19  country 3  county 4  city 6   street 9      1
20  country 3  county 4  city 6        NaN      3
21  country 3  county 4     NaN        NaN      4
22  country 3       NaN     NaN        NaN      5

With a small modification, you can get an output like you are asking for:

dfs_result= list()
group_key= list()
cols_left= cols
index_cols= [f'i_{col}' for col in cols]
while cols_left:
    count_key, *cols_left= cols_left
    group_key.append(count_key)
    df_counts= df.groupby(group_key).agg(count=(count_key, 'count'))
    df_counts.reset_index(drop=False, inplace=True)
    # create temporary index columns and a level column
    df_counts['level']= len(group_key)
    for col in cols:
        if col not in df_counts.columns:
            # create a fake column with NAs, so we have the same structure for
            # all count dfs
            df_counts[f'i_{col}']= ''
        else:
            df_counts[f'i_{col}']= df_counts[col]
        if col != count_key:
            df_counts[col]= ''
    # set the new index (you could also just sort at the end
    # for these columns if you like that better
    df_counts.set_index(index_cols + ['level'], inplace=True)
    dfs_result.append(df_counts)

# concat the partial counts, sort and then drop the 
# temporary index 
df_result= pd.concat(dfs_result, axis='index')
df_result.sort_index(inplace=True, axis='index')
df_result.reset_index(drop=True, inplace=True)
df_result[cols + ['count']]

The final result looks like:

      country    county    city     street  count
0   country 1                                   3
1              county 1                         3
2                        city 1                 2
3                                 street 1      1
4                                 street 2      1
5                        city 2                 1
6                                 street 3      1
7   country 2                                   2
8              county 2                         2
9                        city 3                 2
10                                street 4      1
11                                street 5      1
12  country 3                                   5
13             county 3                         1
14                       city 4                 1
15                                street 6      1
16             county 4                         4
17                       city 5                 1
18                                street 7      1
19                       city 6                 3
20                               street 10      1
21                                street 8      1
22                                street 9      1
Related