How can I move all zeroes to end of array

1

I have a variable that looks like
Input: [[0,2,3,0,0,0,1,2],
[0,0,1,3,0,0,5,0]]

and I want to end up with the following variable
Output: [[2,3,1,2,0,0,0,0],
[1,3,5,0,0,0,0,0]]

I have tried the a[a.nonzero()] if a is the input, but it only works in 1xn dims
Is there a way to do this in PyTorch?

2

You can do it this way:

import torch
x = torch.Tensor([[0,2,3,0,0,0,1,2], [0,0,1,3,0,0,5,0]])
y = torch.empty(0, x.size(1))
for r in x:
    nz = r.nonzero().squeeze()
    z = torch.zeros(r.numel() - nz.numel())
    z = torch.cat((r[nz], z)).unsqueeze(0)
    y = torch.cat((y, z))

print(y)

I’m not sure it’s the best way to do it though, but hopefully it could be good enough for your case.

4

x = torch.Tensor([[0,2,3,0,0,0,1,2], [0,0,1,3,0,0,5,0]])
pad = 0

indices = torch.arange(x.shape[1]).unsqueeze(0).repeat(x.shape[0], 1)
indices[x == pad] = x.shape[1] # set to a BIG number

indices = torch.argsort(indices)
y = torch.gather(x, 1, indices)

5

Masks are a good approach for this. Here is a simple definition for two dimensions, as in your original example:

import torch

def zeros2end(x):
    zeros_mask = x == 0
    out = torch.stack(
        [torch.cat([x[_, :][~zeros_mask[_, :]], x[_, :][zeros_mask[_, :]]], dim=0) for _ in range(x.size()[0])])
    return out

And a test:

torch.manual_seed(0)
x = torch.randint(3, (2, 10))
print("Before")
print(x)
print("After")
print(zeros2end(x))

Gives back:

Before
tensor([[2, 0, 2, 0, 1, 0, 1, 1, 1, 0],
        [2, 2, 0, 0, 1, 2, 0, 0, 2, 0]])
After
tensor([[2, 2, 1, 1, 1, 1, 0, 0, 0, 0],
        [2, 2, 1, 2, 2, 0, 0, 0, 0, 0]])
6 
x = torch.Tensor([[0,2,3,0,0,0,1,2], [0,0,1,3,0,0,5,0]])
print(torch.gather(x, 1, x.ne(0).argsort(dim=1, descending=True, stable=True)))