6
$\begingroup$

I have a sorted list of floats of length n and an n x n list of lists, where each nested list is are also sorted. I want to know the index of the element in list that is nearest to but less than each entry in matrix (essentially, the index in list after which to insert each element of matrix). For example, if n=5:

list=Sort@RandomReal[{-7,7},5]
matrix=Sort/@RandomReal[{-7,7},{5,5}]

For what I’m trying to do, n will be very large (on the order of 2000) and I will have to iterate over this many times, so the code needs to be as fast as possible. Right now, my method for doing this is to flatten the matrix (I won’t actually care about the matrix form in the end---all I want is the index for each entry) and using the following code:

OrdinalFunction[list1_,nearestFn_][list2_]:=With[{r=Flatten@nearestFn[list2]},r-UnitStep[list1[[r]]-list2]]
indexFn=OrdinalFunction[list,Nearest[list->"Index"]];

matrixElements=Flatten@matrix
indexFn[matrixElements]

The advantage of this code is that I can feed the entirety of matrixElements to indexFn to get the indices in list that are nearest to but less than each element. But the code makes no use of the fact that each list in matrix is sorted, and I feel that there must be a faster way to do this. Also, since the lists in matrix are sorted, once one of the elements returns the index of the last elements in ```list```, all subsequent elements will also return this index. Is there a way to take advantage of these facts to make it (ideally much) faster?

$\endgroup$

1 Answer 1

7
$\begingroup$

Suppose that a and b are two sorted lists of real numbers of length n and m, respectively. The following CompiledFunction finds for each element b[[j]] in the second list the desired position in the first list:

cf = Compile[{{a, _Real, 1}, {b, _Real, 1}},
   Block[{i, bj, idx, n, m},
    n = Length[a];
    m = Length[b];
    idx = Table[0, m];
    i = 1;
    Do[
     bj = Compile`GetElement[b, j];
     While[(i <= n) && (Compile`GetElement[a, i] < bj),
      ++i;
      ];
     idx[[j]] = i - 1;
     , {j, 1, m}];
    idx
    ],
   CompilationTarget -> "C",
   RuntimeAttributes -> {Listable},
   Parallelization -> True,
   RuntimeOptions -> "Speed"
   ];

It does so by at most $O(n+m)$ work. Algorithmically, it is similar to the merge step in merge sort.

I am not entirely sure what Nearest[a->"Index"][b] does in this case, but I believe it orders a (because it cannot know that it was ordered before) and then simply performs binary search in the array a. That would take $O( \log_2(n) \, m)$. So the algorithm in cf should be more efficient for large values of n. ( n = 2000) is actually rather small. A stronger advantage of cf is that both arrays are accessed in order, while binary search jumps around in a. So cf should

  1. produce fewer cache misses (this is probably also irrelevant because a fits completely into L1 chache) and
  2. be more predictable for the processor's branch prediction.

Plus cf is parallelized (but Nearest ist probably, too).

This is how the two algorithms compare on my 8-Core machine:

OrdinalFunction[list1_, nearestFn_][list2_] :=  With[{r = Flatten@nearestFn[list2]}, r - UnitStep[list1[[r]] - list2]];
indexFn = OrdinalFunction[list, Nearest[list -> "Index"]];
result = indexFn[Flatten@matrix]; // RepeatedTiming // First
result2 = cf[list, matrix]; // RepeatedTiming // First
ArrayReshape[result, {n, n}] == result2

0.0879834

0.00599676

True

Hence, while Nearest requires $O(n^2 \, \log_2(n))$ work, cf should do all the job in $O(n ( n + n) ) = O(n^2)$. Things change when list has length m independent of n. Then Nearest requires $O(n^2 \, \log_2(m))$ and cf requires $O(n ( m + n) )$ work. The latter grows linearly in m, so Nearest should perform better, when m is several orders of magnitude greater than n.

Edit: Swapped the order of the checks i <= n and Compile`GetElement[a, i] < bj to prevent out-of-bounds reading of a.

$\endgroup$
5
  • $\begingroup$ Amazing! Just to clarify what the code is doing: when you write cf[list,matrix], cf splits matrix into each of its sorted lists because of the specification RuntimeAttributes->{Listable}, correct? Also how could I write a similar code if I also wanted it to work for a sorted list and a matrix whose lists were reverse sorted? Like {7,5.1,0.2,-3,...}. $\endgroup$
    – az123p
    Aug 18, 2022 at 16:46
  • $\begingroup$ Yes, it's exactly the Listable attribute that allows to "thread" through the rows of matrix; and because of Parallelization -> True, it does so in parallel. If the rows of matrix are reverse sorted, then you need only to reverse the Do loop, i.e., Do[ ... , {i,m,1,-1}];. $\endgroup$ Aug 18, 2022 at 16:56
  • $\begingroup$ Perfect! Thanks for the help. $\endgroup$
    – az123p
    Aug 18, 2022 at 17:08
  • $\begingroup$ You're welcome! $\endgroup$ Aug 18, 2022 at 17:08
  • $\begingroup$ Quick follow-up: since we have Parallelization->True inside the compiled function, what would happen if I used cf in a piece of code that is inside, say, ParallelTable? I assume that there would be some kind of conflict over resources. Would Parallelization automatically be reset to False in cf, or would there still be a benefit to keeping this piece in? $\endgroup$
    – az123p
    Aug 22, 2022 at 6:37

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.