Slice a 2D array using built-in slice function

I need to write a function that returns a single row or a column from a 2D array. The input to the function tells what to return.

# 3x3 array with values 1 to 9
a = np.arange(1, 10).reshape(3,3)
rowCount, colCount = a.shape

# return last row [7, 8, 9]
a[rowCount - 1, :]

# return first row [1, 2, 3]
a[0, :]

# return last column [3, 6, 9]
a[:, lastCol]

# return first column [1, 4, 7]
a[:, 0]

How can I do this in a function, such that the function receives the row or column to return?

Something like,

def getSlice(slice):
    return a[slice]

Where the slice object is created using the built-in slice function.

However, I cannot figure out how to create a slice object for a 2D array, due to the fact that slice function does not accept the colon operator like slice(0, :).

Also, is there a way to represent "last row" or "last column if I do not know beforehand the shape of the 2D array?

Use-Case

Below are a couple of use-cases why I need a function instead of directly using the a[:, 0] expression:

1. The caller does not have access to the array. The caller can get a desired row or column from the array by calling the getSlice function.
2. The preferred row or column needs to be pre-configured. For instance, {a1: 'first row', a2: 'last column'}. Both a1 and a2 may get transposed and modified many times. But at all times, I am interested only in the configured row/column of the two arrays.

I will try to pick your question apart a bit:

> I cannot figure out how to create a slice object for a 2D array, due to the fact that slice function does not accept the colon operator like slice(0, :)

Indeed, this does not work, as slice objects do not work like that: they represent the index range along one dimension, while you want to index along two dimensions. So, while this does not work with a single slice object, you can use a tuple – compare the doc here, especially the last bullet point of the section Slicing and striding:

> A slicing tuple can always be constructed as obj and used in the x[obj] notation. Slice objects can be used in the construction in place of the [start:stop:step] notation. For example, x[1:10:5, ::-1] can also be implemented as obj = (slice(1, 10, 5), slice(None, None, -1)); x[obj].

In your, case, as you only ask for complete rows and columns:

> I need to write a function that returns a single row or a column from a 2D array.

… the corresponding : can be written as slice(None, None, None). In fact, you do not even need to explicitly create the slicing tuple, as you only ever will exactly have two indices (one for row, one for column, either of them being :) that you can directly use for indexing your array. So, if you really want to pack your indexing into a function, I would create a function as follows:

def get_slice(row=slice(None,None,None), col=slice(None,None,None)):
    return a[row, col]  # or, equivalently a[(row, col)]

Here, we will provide :, or rather, the equivalent slice(None,None,None), as the default value for both row and column, so that the users only need to pass one value for the index that they are actually interested in. Of course, you can also adjust the function so that a tuple is accepted instead, which would be closer to your original code.

In any case, the version above will give you:

print(get_slice(row=0))  # First row
# >>> [1 2 3]
print(get_slice(col=-1))  # Last column
# >>> [3 6 9]
print(get_slice(col=0))  # First column
# >>> [1 4 7]

Finally:
> Also, is there a way to represent "last row" or "last column" …?

You can achieve this with negative indexing, which I already used for the last column in the example above. Negative indexing works "from the right" or "from the bottom", thus a[-1] refers to the last row, a[:, -2] refers to the penultimate column, etc., and does not require knowing the actual shape of the array.

While a is: [[1 2 3][4 5 6][7 8 9]].

a[0] -> [1 2 3]
a[1] -> [4 5 6]
a[2] -> [7 8 9]

if you wan't 4 ->

a[1][0] -> 4

So if you treat is as matrix then the function can accept 2 parameters (x, y)

def get_element(matrix, x, y):
    return matrix[x][y]

If you wan't column - that there I suggest to apply separated method - due to Separation of Concerns principle:

def select_column(matrix, column_index):
    selected_column = [row[column_index] for row in matrix]
    return selected_column

test = select_column(a, 1) -> [2, 5, 8]

or in more concise manner:

a[:,1]