Compute Arithmetic Mean Along Matrix Axis

In data analysis and machine learning, computing statistical aggregations along specific dimensions of a matrix is a foundational operation. One of the most common aggregations is the arithmetic mean, which provides the central tendency of a set of values.

Given a matrix M of dimensions n × m (where n is the number of rows and m is the number of columns), you can compute the mean along two different axes:

Row-wise Mean (mode = 'row'): For each row i, compute the average of all elements in that row: $$\text{mean}i = \frac{1}{m} \sum{j=1}^{m} M_{ij}$$ This produces a vector of n mean values, one for each row.

Column-wise Mean (mode = 'column'): For each column j, compute the average of all elements in that column: $$\text{mean}j = \frac{1}{n} \sum{i=1}^{n} M_{ij}$$ This produces a vector of m mean values, one for each column.

Intuition:

• Row-wise aggregation summarizes each individual observation across all its features
• Column-wise aggregation summarizes each feature across all observations

Your Task: Write a Python function that computes the arithmetic mean of a matrix along a specified axis. The function should accept a matrix (2D list) and a mode string ('row' or 'column'), and return a list of mean values corresponding to the specified aggregation direction.

matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
mode = 'column'

[4.0, 5.0, 6.0]

matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
mode = 'row'

[2.0, 5.0, 8.0]

matrix = [[1, 2], [3, 4]]
mode = 'column'

[2.0, 3.0]