Comment obtenir l'indice de la colonne à partir de son nom avec pandas ?

Exemples de comment obtenir l'indice de la colonne à partir de son nom avec pandas

Créez un dataframe avec pandas

Commençons par créer un dataframe simple avec pandas:

import pandas as pd
import numpy as np

data = np.random.randint(10, size=(5,3))

columns = ['Score A','Score B','Score C']

df = pd.DataFrame(data=data,columns=columns)

print(df)

donne par exemple

   Score A  Score B  Score C
0        3        5        6
1        8        1        5
2        3        5        3
3        3        4        6
4        4        6        7

Obtenir l'indice de la colonne à partir de son nom (le nom de la colonne est unique)

Pour obtenir l'indice de la colonne à partir de son nom, une solution est d'utiliser get_loc (), exemple:

df.columns.get_loc("Score A")

donne

0

Colonne intitulée "Score B"

df.columns.get_loc("Score B")

donne

1

Colonne intitulée "Score C"

df.columns.get_loc("Score C")

donne

2

Il est possible de créer un dictionnaire avec le nom de la colonne comme clé et l'indice associé comme valeur:

idx_dic = {}

for col in df.columns:
    idx_dic[col] = df.columns.get_loc(col)

print(idx_dic)

donne

{'Score A': 0, 'Score B': 1, 'Score C': 2}

Cas 2: plusieurs colonnes adjacentes avec le même nom

Un autre cas, ici il y a plusieurs colonnes adjacentes avec le meme nom "Score C"

columns = ['Score A','Score B','Score C']
df = pd.DataFrame(data=data,columns=columns)

data2 = np.random.randint(10, size=(5,1))

columns = ['Score C']

df_add = pd.DataFrame(data=data2,columns=columns)

df = pd.concat([df,df_add], axis=1)
df = pd.concat([df,df_add], axis=1)

print(df)

donne

   Score A  Score B  Score C  Score C  Score C
0        3        5        6        3        3
1        8        1        5        1        1
2        3        5        3        3        3
3        3        4        6        2        2
4        4        6        7        7        7

alors

df.columns.get_loc("Score C")

donne un objet de type slice:

slice(2, 5, None)

Rappel: pour obtenir le type

type(df.columns.get_loc("Score C"))

autre approche

isinstance(df.columns.get_loc("Score C"), slice)

donne ici

True

Donc pour obtenir une liste d'index associée à la colonne "Score C", une solution est de faire:

col_idx_slice = df.columns.get_loc("Score C")

col_idx_list = [i+col_idx_slice.start for i in range(col_idx_slice.stop-col_idx_slice.start) ]

col_idx_list

[2, 3, 4]

Cas 3: plusieurs colonnes non adjacentes avec le même nom

Autre cas possible plusieurs colonnes non adjacentes avec le même nom "Score C":

columns = ['Score A','Score B','Score C']
df = pd.DataFrame(data=data,columns=columns)

data2 = np.random.randint(10, size=(5,2))

columns = ['Score C','Score D']

df_add = pd.DataFrame(data=data2,columns=columns)

df = pd.concat([df,df_add], axis=1)
df = pd.concat([df,df_add], axis=1)

print(df)

donne

   Score A  Score B  Score C  Score C  Score D  Score C  Score D
0        3        5        6        3        9        3        9
1        8        1        5        0        4        0        4
2        3        5        3        3        0        3        0
3        3        4        6        6        9        6        9
4        4        6        7        7        4        7        4

alors

df.columns.get_loc("Score C")

donne ici une matrice

array([False, False,  True,  True, False,  True, False])

Note: pour retrouver le type:

type(df.columns.get_loc("Score C"))

numpy.ndarray

ou

isinstance(df.columns.get_loc("Score C"), np.ndarray)

donne ici

True

Pour convertir ce tableau booléen en indices:

np.where( df.columns.get_loc("Score C") )[0]

ce qui donne

array([2, 3, 5])

et ensuite en liste:

np.where( df.columns.get_loc("Score C") )[0].tolist()

donne

[2, 3, 5]

Créer une fonction pour obtenir des index à partir du nom de la colonne

Combinez les 3 exemples ci-dessus pour créer une fonction qui renvoie l'indice(s) de colonne à partir de son nom:

Exemple

   Score A  Score B  Score C  Score D  Score D  Score C
0        3        5        6        6        7        5
1        8        1        5        7        3        9
2        3        5        3        0        1        7
3        3        4        6        8        2        2
4        4        6        7        5        0        1

Solution

def get_column_index_list(column):

    if type( df.columns.get_loc(column) ) == int:
        col_index_list = []
        col_index_list.append( df.columns.get_loc(column) )
    if isinstance(df.columns.get_loc(column), slice):
        col_index_list = [i+col_idx_slice.start for i in range(col_idx_slice.stop-col_idx_slice.start) ]
    if isinstance(df.columns.get_loc(column), np.ndarray): 
        col_index_list = np.where( df.columns.get_loc(column) )[0].tolist()

    return col_index_list

print( get_column_index_list('Score A') )
print( get_column_index_list('Score C') )
print( get_column_index_list('Score D') )

donne

[0]
[2, 5]
[3, 4]

Références

• Get column index from column name of a given Pandas DataFrame
• Get column index from column name in python pandas
• pandas.Index.get_loc