Fastest way to (arg)sort a flattened nD-array that is sorted along each dimension?

user_6753182 0
Viewed 240

The question itself is language-agnostic. I will use python for my example, mainly because I think it is nice to demonstrate the point.

I have an N-dimensional array of shape (n1, n2, ..., nN) that is contiguous in memory (c-order) and filled with numbers. For each dimension by itself, the numbers are ordered in ascending order. A 2D example of such an array is:

>>> import numpy as np
>>> n1 = np.arange(5)[:, None]
>>> n2 = np.arange(7)[None, :]
>>> n1+n2
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 1,  2,  3,  4,  5,  6,  7],
       [ 2,  3,  4,  5,  6,  7,  8],
       [ 3,  4,  5,  6,  7,  8,  9],
       [ 4,  5,  6,  7,  8,  9, 10]])

In this case, the values in each row are ascending, and the values in each column are ascending, too. A 1D example array is

>>> n1 = np.arange(10)
>>> n1*n1
array([ 0,  1,  4,  9, 16, 25, 36, 49, 64, 81])

I would like to obtain a list/array containing the indices that would sort the flattened version of the nD array in ascending order. By the flattened array I mean that I interpret the nD-array as a 1D array of equivalent size. The sorting doesn't have to preserve order, i.e., the order of indices indexing equal numbers doesn't matter. For example

>>> n1 = np.arange(5)[:, None]
>>> n2 = np.arange(7)[None, :]
>>> arr = n1*n2
>>> arr
array([[ 0,  0,  0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4,  5,  6],
       [ 0,  2,  4,  6,  8, 10, 12],
       [ 0,  3,  6,  9, 12, 15, 18],
       [ 0,  4,  8, 12, 16, 20, 24]])
>>> np.argsort(arr.ravel())
array([ 0, 28, 14,  7,  6, 21,  4,  3,  2,  1,  5,  8,  9, 15, 22, 10, 11,
       29, 16, 12, 23, 17, 13, 18, 30, 24, 19, 25, 31, 20, 26, 32, 27, 33,
       34], dtype=int64)

Standard sorting on the flattened array can accomplish this; however, it doesn't exploit the fact that the array is already partially sorted, so I suspect there exists a more efficient solution. What is the most efficient way to do so?

A comment asked what my use-case is, and if I could provide some more realistic test data for benchmarking. Here is how I encountered this problem:

Given an image and a binary mask for that image (which selects pixels), find the largest sub-image which contains only selected pixels.

In my case, I applied a perspective transformation to an image, and want to crop it so that there is no black background while preserving as much of the image as possible.

from skimage import data
from skimage import transform
from skimage import img_as_float

tform = transform.EuclideanTransform(
    rotation=np.pi / 12.,
    translation = (10, -10)
    )

img = img_as_float(data.chelsea())[50:100, 150:200]
tf_img = transform.warp(img, tform.inverse)
tf_mask = transform.warp(np.ones_like(img), tform.inverse)[..., 0]

y = np.arange(tf_mask.shape[0])
x = np.arange(tf_mask.shape[1])
y1 = y[:, None, None, None]
y2 = y[None, None, :, None]
x1 = x[None, :, None, None]
x2 = x[None, None, None, :]

y_padded, x_padded = np.where(tf_mask==0.0)
y_padded = y_padded[None, None, None, None, :]
x_padded = x_padded[None, None, None, None, :]
y_inside = np.logical_and(y1[..., None] <= y_padded, y_padded<= y2[..., None])
x_inside = np.logical_and(x1[..., None] <= x_padded, x_padded<= x2[..., None])
contains_padding = np.any(np.logical_and(y_inside, x_inside), axis=-1)

# size of the sub-image
height = np.clip(y2 - y1 + 1, 0, None)
width = np.clip(x2 - x1 + 1, 0, None)
img_size = width * height

# find all largest sub-images
img_size[contains_padding] = 0
y_low, x_low, y_high, x_high = np.where(img_size == np.max(img_size))
cropped_img = tf_img[y_low[0]:y_high[0]+1, x_low[0]:x_high[0]+1]

The algorithm is quite inefficient; I am aware. What is interesting for this question is img_size, which is a (50,50,50,50) 4D-array that is ordered as described above. Currently I do:

img_size[contains_padding] = 0
y_low, x_low, y_high, x_high = np.where(img_size == np.max(img_size))

but with a proper argsort algorithm (that I can interrupt early) this could potentially be made much better.

language-agnostic sorting arrays
2 Answers

I would do it using parts of mergesort and a divide and conquer approach. You start with the first two arrays.

[0,  1,  2,  3,  4,  5,  6],//<- This
[ 1,  2,  3,  4,  5,  6,  7],//<- This
....

Then you can merge them like this (Java-like syntax):

List<Integer> merged=new ArrayList<>();
List<Integer> firstRow=... //Same would work with arrays
List<Integer> secondRow=...
int firstCnter=0;
int secondCnter=0;
while(firstCnter<firstRow.size()||secondCnter<secondRow.size()){
  if(firstCnter==firstRow.size()){ //Unconditionally add all elements from the second, if we added all the elements from the first
     merged.add(secondRow.get(secondCnter++)); 
  }else if(secondCnter==secondRow.size()){
     merged.add(firstRow.get(firstCnter++));
  }else{ //Add the smaller value from both lists at the current index.
    int firstValue=firstRow.get(firstCnter);
    int secondValue=secondRow.get(secondCnter);
    merged.add(Math.min(firstValue,secondValue));
    if(firstValue<=secondValue)
       firstCnter++;
    else
       secondCnter++;
  }
}

After that you can merge the next two rows, until you have:

[0,1,1,2,2,3,3,4,4,5,5,6,7]
[2,3,3,4,4,5,5,6,6,7,7,8,8,9]
[4,5,6,7,8,9,10] //Not merged.

Continue to merge again.

[0,1,1,2,2,2,3,3,3,4,4,4,4,5,5,5,6,6,6,7,7,7,8,8,9]
[4,5,6,7,8,9,10]

After that, the last merge:

[0,1,1,2,2,2,3,3,3,4,4,4,4,4,5,5,5,5,6,6,6,6,7,7,7,7,8,8,8,9,9,10]

I don't know about the time complexity, but should be a viable solution

user_139121320

Another idea: Use a min-heap with just the current must-have candidates for being the next-smallest value. Start with the value at the origin (index 0 in all dimensions), as that's smallest. Then repeatedly take out the smallest value from the heap and add its neighbors not yet added.

user_130084390
Related