5 Best Ways to Retrieve Category Codes in Python Pandas CategoricalIndex

by

💡 Problem Formulation: Working with categorical data in pandas often involves converting textual data into a categorical type for efficiency and ease of analysis. Sometimes, we need to retrieve the integer codes of categories from a CategoricalIndex. This article illustrates how you can extract the underlying category codes from a pandas CategoricalIndex object, with an input being the categorical object and the output being the corresponding codes for each category.

Method 1: Using categories.codes Attribute

Retrieving category codes in pandas is straightforward using the .codes attribute of the categories object within a CategoricalIndex. This attribute returns an array of integers coding the categories (provided the dtype is categorical).

Here’s an example:

import pandas as pd

# Create a categorical series
categorical_series = pd.Series(["red", "blue", "red", "green"]).astype('category')

# Access the category codes
codes = categorical_series.cat.codes
print(codes)

Output:

0    2
1    0
2    2
3    1
dtype: int8

Here, the categorical_series is converted to a categorical type, and cat.codes is used to extract the category codes, where each integer represents a specific category as determined by pandas during the conversion process.

Method 2: Using CategoricalIndex.codes

The CategoricalIndex also has a .codes property similar to the one used in method 1. If your dataset index is of type CategoricalIndex, you can directly access the codes in this manner without an additional conversion step.

Here’s an example:

import pandas as pd

# Create a categorical index
categories = ["small", "medium", "large"]
categorical_index = pd.CategoricalIndex(categories)

# Get category codes
category_codes = categorical_index.codes
print(category_codes)

Output:

[2 1 0]

This code snippet demonstrates obtaining category codes from a CategoricalIndex object. The output reflects the codes assigned to “small”, “medium”, and “large” following pandas’ internal ordering logic.

Method 3: Using get_codes() function

Although not a common approach, one can also use a custom function like get_codes() to iterate over the categorical data and manually map each category to its corresponding code. This is particularly useful if further manipulation of the codes is needed during retrieval.

Here’s an example:

import pandas as pd

# Create a categorical series
categorical_series = pd.Series(["apple", "orange", "apple", "banana"]).astype('category')

# Define get_codes function
def get_codes(series):
    return series.cat.codes

# Apply the function
codes = get_codes(categorical_series)
print(codes)

Output:

0    0
1    2
2    0
3    1
dtype: int8

This approach introduces a function, get_codes(), which returns the category codes for a given series. This is primarily a reusability enhancement, should codes be needed across multiple series in a consistent manner.

Method 4: Using apply() Method

For more complex data manipulations involving category codes, pandas’ apply() method can be employed to run a function across the series’ values. This method is highly versatile and can accommodate additional logic within the function being applied.

Here’s an example:

import pandas as pd

# Create a categorical dataframe
df = pd.DataFrame({'fruits': ["apple", "orange", "apple", "banana"]}).astype('category')

# Apply custom function to get codes
codes = df['fruits'].apply(lambda x: x.cat.codes)
print(codes)

Output:

0    0
1    2
2    0
3    1
Name: fruits, dtype: int8

This code uses apply() with a lambda function to extract the codes for each category in a DataFrame column. It is especially useful when working with a DataFrame rather than a series.

Bonus One-Liner Method 5: List Comprehension

Sometimes simplicity is key. List comprehension can be a one-liner solution for extracting the codes from a categorical series, especially when inline operations are required.

Here’s an example:

import pandas as pd

# Create a categorical series
categorical_series = pd.Series(["spring", "summer", "fall", "winter"]).astype('category')

# List comprehension to get codes
codes = [code for code in categorical_series.cat.codes]
print(codes)

Output:

[3, 2, 0, 1]

In this single line of code, list comprehension iterates over the cat.codes attribute of the categorical series and constructs a list of the category codes.

Summary/Discussion

  • Method 1: Use of .codes attribute. Straightforward and the most common way to get codes of categorical data. Limited flexibility for complex operations.
  • Method 2: Direct .codes property on the CategoricalIndex. Best for when working directly with a CategoricalIndex, bypassing the need for series or dataframe manipulations.
  • Method 3: Creation of a custom get_codes() function. Offers modularity and is best suited for use-cases requiring the same logic in multiple places. Potentially overkill for simple code extraction.
  • Method 4: Use of apply() method. Flexible and powerful, ideal for complex operations with additional logic required. Might be less performance-efficient compared to direct attribute access.
  • Bonus Method 5: List comprehension. Quick and concise, perfect for inline operations or scripts with minimal complexity. However, it may not be the most readable approach for those new to Python.