Slicing 2D list (inner dimension) in Python and corresponding method in Numpy array and PyTorch tensor

Objective

There are a lot of time I slice some portion of data from multi-dimension vector/tensor. The Numpy array and PyTorch tensor make it very easy to slice, and with a very similar syntax. In some scenario I might need to work from a list, and here comes one implementation that can be done.

Example 1

A simple 2D number list, I want to slice the input into 3 list like elements

from:
[[ 1, 2, 3, 4, 5],
 [11,12,13,14,15],
 [21,22,23,24,25]]
into:
[[ 1, 2, 3], | [4, | [7,
 [11,12,13], |  5, |  8,
 [21,22,23]] |  6] |  9]

Numpy array

The syntax is using [: , :x] and [: , y], this means for the first dimension, slice all, and the 2nd dimension, slice :x list or slice yth item

PyTorch tensor

Same slicing syntax as Numpy

List

Using a zip, * and list comprehension, this does not look nature.

The innermost list comprehension, it return the result:

[([1, 2, 3], 4, 5),
 ([11, 12, 13], 14, 15),
 ([21, 22, 23], 24, 25)]

Explanation of the * => zip

// With the *, it spread the innermost list comprehension return
[([1, 2, 3], 4, 5),
 ([11, 12, 13], 14, 15),
 ([21, 22, 23], 24, 25)]
// into 3 items (I think is as tuple because one cannot really * and check value)
(([1, 2, 3], 4, 5),
 ([11, 12, 13], 14, 15),
 ([21, 22, 23], 24, 25))
// with the zip operation, it rearrange the result as I wished
[([1, 2, 3], [11, 12, 13], [21, 22, 23]),  
 (4, 14, 24),
 (5, 15, 25)]

Example 2

This example, I am having some string information:

Numpy array

Numpy array would do the expected result, even with string

PyTorch tensor

Pytorch tensor is for numerical calculation, so the string input would break it.

List

List operation work as expected.

Final words

The Numpy array do perform very versatile, the list approach is some backup plan for myself, and I didn’t check performance and memory pressure comparing 1) direct list operation vs 2) list => numpy array => list operation.