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.

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