Can numpy do better than this?
Ian Kelly
ian.g.kelly at gmail.com
Thu Jan 8 14:27:46 EST 2015
On Thu, Jan 8, 2015 at 10:56 AM, Rustom Mody <rustompmody at gmail.com> wrote:
> Given a matrix I want to shift the 1st column 0 (ie leave as is)
> 2nd by one place, 3rd by 2 places etc.
>
> This code works.
> But I wonder if numpy can do it shorter and simpler.
>
> ---------------------
> def transpose(mat):
> return([[l[i] for l in mat]for i in range(0,len(mat[0]))])
> def rotate(mat):
> return([mat[i][i:]+mat[i][:i] for i in range(0, len(mat))])
> def shiftcols(mat):
> return ( transpose(rotate(transpose(mat))))
Without using numpy, your transpose function could be:
def transpose(mat):
return list(zip(*mat))
numpy provides the roll function, but it doesn't allow for a varying
shift per index. I don't see a way to do it other than to roll each
column separately:
>>> mat = np.array([[1,2,3,4,5,6],
... [7,8,9,10,11,12],
... [13,14,15,16,17,18],
... [19,20,21,22,23,24],
... [25,26,27,28,29,30],
... [31,32,33,34,35,36],
... [37,38,39,40,41,42]])
>>> res = np.empty_like(mat)
>>> for i in range(mat.shape[1]):
... res[:,i] = np.roll(mat[:,i], -i, 0)
...
>>> res
array([[ 1, 8, 15, 22, 29, 36],
[ 7, 14, 21, 28, 35, 42],
[13, 20, 27, 34, 41, 6],
[19, 26, 33, 40, 5, 12],
[25, 32, 39, 4, 11, 18],
[31, 38, 3, 10, 17, 24],
[37, 2, 9, 16, 23, 30]])
More information about the Python-list
mailing list