Argmax along axis

Implement a function that returns the index of the maximum value along a specified axis of a 2D list.

lists max index nested-loops

Write a function `argmax_axis(arr, axis)` that takes a 2D list of numbers `arr` (non-empty, rectangular) and an integer `axis` (0 or 1) and returns a list of integers representing the indices of the maximum values along the specified axis.

- If `axis=0`, for each column, return the row index of the maximum value in that column.
- If `axis=1`, for each row, return the column index of the maximum value in that row.

If there are multiple occurrences of a maximum in a row/column, return the smallest index (first occurrence). The input will always be a 2D list with at least one element, and all rows will have the same length. The output must be a list of integers.

Constraints

- `arr` is a 2D list of numbers (ints or floats) with shape (m, n), where m ≥ 1 and n ≥ 1.
- All rows have equal length.
- `axis` is either 0 or 1.
- Numbers may be negative.
- Time complexity: O(m*n) per call.

Example

```python
# Example 1:
arr = [[5, 2, 9],
       [1, 8, 3],
       [7, 4, 6]]
argmax_axis(arr, 0)
# Output: [2, 1, 0]

argmax_axis(arr, 1)
# Output: [2, 1, 0]

# Example 2:
arr2 = [[1, 1, 1],
        [2, 2, 2]]
argmax_axis(arr2, 0)
# Output: [1, 1, 1]

argmax_axis(arr2, 1)
# Output: [0, 0]
```

10 points ~10 min

Recent Submissions

No submissions yet — hit Run Tests to try!

Hints

Think of a loop over rows or columns, depending on the axis.

For axis=1, each row is a list: find the max and its first index.

For axis=0, you need to scan down each column: track the current max and its row index.

Start with the first element as the initial max index, then update when you find a strictly larger value.

Test Results

Press Ctrl+Enter or click Run Tests to execute your code.