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]

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

Same slicing syntax as Numpy

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 would do the expected result, even with string

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

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.

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store