Sorting the contents of the matrix and then picking the index with the median value is a good approach. Lets see if we can do it with constant extra memory.
for i in range(1,len(A)):
This extends the first row of the matrix to contain every row in a flat list. Before the matrix was of size N * M, whereas now it is N * M (the first row + (N - 1) * M (all the other rows). Subtracting the original size from this tells us how much extra memory we are using. We use (N - 1) * M additional memory or in other words O(NM) extra memory. This is not what we want.
The reason to put all the elements in one list is to make sorting easy. Lets see if we can sort without needing a flatten (1d) list. There are many sorts that don't require extra memory, they are called "inplace" sorting algorithms. For simplicity we will modify selection sort to work for our case.
How selection sort works is it picks the smallest element in the list, and puts it at the front. Then it finds the next smallest element, and puts it second, and so forth. To implement this, we can find the smallest in the whole list, and swap it with the first element. Then we can find the smallest of the list skipping the first slot.
def index_of_smallest(numbers, starting_index):
# Assume numbers is not empty.
smallest, index = numbers[starting_index], starting_index
for i, number in enumerate(numbers[starting_index:], starting_index):
if number < smallest:
smallest, index = number, i
size = len(numbers)
for i in range(size):
index = index_of_smallest(numbers, i)
numbers[i], numbers[index] = numbers[index], numbers[i]
# Don't return anything, we are modifying it inplace.
Now, we need this process to work on a matrix instead of a flat list. This is straightforward enough, we can loop over the matrix (left to right, top to bottom) and ignore cells we have already dealt with. In the below code x is the row coordinate, and y is the column coordinate.
def coordinates_of_smallest(matrix, starting_x, starting_y):
smallest, smallest_x, smallest_y = matrix[starting_x][starting_y], starting_x, starting_y
for x, row in enumerate(matrix):
for y, cell in enumerate(row):
if x < starting_x or (x == starting_x and y < starting_y):
if cell < smallest:
smallest, smallest_x, smallest_y = cell, x, y
return smallest_x, smallest_y
# Assume the matrix is not empty.
n, m = len(matrix), len(matrix)
for x in range(n):
for y in range(m):
smallest_x, smallest_y = coordinates_of_smallest(matrix, x, y)
matrix[x][y], matrix[smallest_x][smallest_y] = matrix[smallest_x][smallest_y], matrix[x][y]
>>> matrix = [[1, 3, 5], [2, 6, 9], [3, 6, 9]]
>>> print(matrix) # [[1, 2, 3], [3, 5, 6], [6, 9, 9]]
Now getting the median of this is a piece of cake, it will be in the middle slot of the middle row! Since N * M is odd, both N and M must be odd. Therefore the median is at matrix[N // 2][M // 2].
There is a little room for improvement here. While we only use constant extra memory, our time complexity has gone up from O(nm lognm) to O((nm)**2). For a better time complexity, I would recommend using inplace quicksort which brings us back to O(nm lognm).
Another point is that we are doing too much work. Once we have worked our way up to the row N // 2 and the slot M // 2, we are actually done! We have put the median element in it's place, and we can stop. This is a simple enough check to add, but can cut the actual running time of the code in half.