Add headers to colums and rows in Python - 2D array (list of lists)

Question

I have a 2D Python array (list of lists). I treat this as a 'table'. Now I want to replace the first row and the top row with a header row and column. The row and the header are exactly the same.

I try this using list comprehension. It works out for the 'row' part, but the column part is not very Pythonic yet.

# The header that i want to add
headers = ['foo', 'bar', 'baz', 'other']
l = len(headers) + 1 # The final matrix is one element bigger
# square matrix of random size, filled with data
array = [['xxx' for i in range(l)] for j in range(l)]
# The concise one - add headers to top row
array[0] = ['Title'] + [category for category in headers]
# The ugly one - add headers for the rows
for i in range(l):
 array[i][0] = array[0][i]

The final output should look like this (which it does):

[['Title', 'foo', 'bar', 'baz', 'other'],
 ['foo', 'xxx', 'xxx', 'xxx', 'xxx'],
 ['bar', 'xxx', 'xxx', 'xxx', 'xxx'],
 ['baz', 'xxx', 'xxx', 'xxx', 'xxx'],
 ['other', 'xxx', 'xxx', 'xxx', 'xxx']]

I'm just not so happy with the 'for' loop. How can this be done more Pythonic?

Answer 1

Maybe this can help (if you don't want to use numpy):

headers = ['foo', 'bar', 'baz', 'other']
l = len(headers)
arr = [["xxx" for i in range(l)] for j in range(l)]
# adding top row
arr = [headers] + arr
# adding first column
headers_mod = ['Title'] + headers
new_arr = [[headers_mod[i]]+arr[i] for i in range(l+1)]
for i in new_arr:
 print(*i)

gives you the output as:

Title foo bar baz other
foo xxx xxx xxx xxx
bar xxx xxx xxx xxx
baz xxx xxx xxx xxx
other xxx xxx xxx xxx

---

Otherwise, when dealing with array manipulations in python try going with numpy, pandas, as they provide better operations like by giving option for axis, transpose, etc.

• grrfield: After rethinking you are right. I'm just not that familiar wit NumPy. I changed one of my matrices to a numpy array and surprisingly my code was still working. Thank you.
• Shyam Mittal: Yeah, mostly its the same but with very useful features for mathematical manipulations.

Answer 2

numpy is excellent for tables, but for a labeled table like this pandas might be better for your needs.

Solution using numpy:

import numpy as np
# The header that i want to add
headers = ['foo', 'bar', 'baz', 'other']
ll = len(headers)+1
data = [['xxx' for _ in range(ll)] for j in range(ll)]
data = np.array(data, dtype=object)
data[0,0] = 'Title'
data[0,1:] = headers
data[1:,0] = headers
print(data)

prints

[['Title' 'foo' 'bar' 'baz' 'other']
 ['foo' 'xxx' 'xxx' 'xxx' 'xxx']
 ['bar' 'xxx' 'xxx' 'xxx' 'xxx']
 ['baz' 'xxx' 'xxx' 'xxx' 'xxx']
 ['other' 'xxx' 'xxx' 'xxx' 'xxx']]

Setting dtype to object allows your array to mix strings and other data types you might want to use. If your data is just strings then you can use 'UN' as the dtype, where N is the longest string you plan to use. (Numpy, when making an all string array automatically picks your longest string as the maximum length for the strings, which is fine unless your strings are all shorter than the headers you plan to add.)

Alternate version of the above code:

import numpy as np
# The header that i want to add
headers = ['foo', 'bar', 'baz', 'other']
# Add Title to headers to simply later assignment
headers = ['Title'] + headers
ll = len(headers)
data = [['xxx' for _ in range(ll)] for j in range(ll)]
data = np.array(data)
data[0,:] = headers
data[:,0] = headers
print(data)

pandas, on the other hand, is explicitly designed to handle headers

import numpy as np, pandas as pd
# The header that i want to add
headers = ['foo', 'bar', 'baz', 'other']
ll = len(headers) + 1
data = [['xxx' for _ in range(ll)] for j in range(ll)]
data = np.array(data)
data = pd.DataFrame(data[1:,1:], columns=headers, index=headers)
data.columns.name = 'Title'
data.loc['foo','bar'] = 'yes'
print(data)
print('')
print(data['bar'])
print('')
print(data.loc['foo',:])

prints

Title foo bar baz other
foo xxx yes xxx xxx
bar xxx xxx xxx xxx
baz xxx xxx xxx xxx
other xxx xxx xxx xxx
foo yes
bar xxx
baz xxx
other xxx
Name: bar, dtype: object
Title
foo xxx
bar yes
baz xxx
other xxx
Name: foo, dtype: object