Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Get column names including index name from pandas data frame

Assuming that we have a data frame with an index that might have a name:

df = pd.DataFrame({'a':[1,2,3],'b':[3,6,1], 'c':[2,6,0]})
df = df.set_index(['a'])

   b  c
a      
1  3  2
2  6  6

What is the best way to get the column names that will include the index name if it is present.

Calling df.columns.tolist() do not include the index name and return ['b', 'c'] in this case, and I would like to obtain ['a', 'b', 'c'].

like image 494
Krzysztof Słowiński Avatar asked Jan 28 '23 18:01

Krzysztof Słowiński


2 Answers

The index can be temporarily reset for the call:

df.reset_index().columns.tolist()

If an empty index name is not to appear in the list, do the reset_index() conditionally:

(df.reset_index() if df.index.name else df).columns.tolist()
like image 57
w-m Avatar answered Jan 30 '23 08:01

w-m


For universal solution need filter None if not exist index.name:

df = pd.DataFrame({'a':[1,2,3],'b':[3,6,1], 'c':[2,6,0]})

print ([df.index.name] + df.columns.tolist())
[None, 'a', 'b', 'c']

c = list(filter(None, [df.index.name] + df.columns.tolist()))
print (c)
['a', 'b', 'c']

df = df.set_index(['a'])

c = list(filter(None, [df.index.name] + df.columns.tolist()))
print (c)
['a', 'b', 'c']

Another solution with numpy.insert and difference:

c = np.insert(df.columns, 0, df.index.name).difference([None]).tolist()
print (c)

['a', 'b', 'c']
like image 29
jezrael Avatar answered Jan 30 '23 06:01

jezrael