Efficient extraction of a subset

Question

The context is molecular dynamics. The data is in the format (index_1, index_2, distance). The goal is to find all of the distances to the neighbors of point index_0 and to present a set listing the neighbor and the distance.

For example, given a data set like:

data = {{1, 2, Subscript[a, 1]}, {1, 3, Subscript[a, 2]}, {1, 4, Subscript[a,3]}, {1, 5, Subscript[a, 4]}, {2, 3, Subscript[a, 5]}, {2, 4, Subscript[a, 6]}, {2, 5, Subscript[a, 7]}, {3, 4, Subscript[a, 8]}, {3, 5, Subscript[a, 9]}, {4, 5, Subscript[a, 10]}}

A desired output for the distances to all of the neighbors of point 2 is:

output = {{1, Subscript[a, 1]}, {3, Subscript[a, 5]}, {4, Subscript[a, 6]}, {5, Subscript[a, 7]}}

1) The distance list data is an input.

2) Efficiency is a premium here as the routine will be used heavily in lieu of storing the lists output. We are trading flops for memory movement.

3) If convenient it would be nice to sort the list from smallest distance to largest distance.

Answer 1

Without further specification of your needs (c.f. Mr. Wizard's comment), here's a fast method:

nf = With[{td = Transpose@data},
    Nearest[
     Join @@ td[[;; 2]] -> 
      Join[Transpose[td[[2 ;;]]], Transpose[td[[{1, 3}]]]]]]; // Timing

results= SortBy[#, Last] & /@ nf /@ Union @@ data[[All, ;; 2]]; // Timing

Cobbling up some data via:

data = DeleteCases[RandomInteger[2000, {2000, 3}], {x_, x_, _}];

this is ~100X faster (including time to create nearest function and sorting results for all in order of distance) vs

Cases[data, {f___, #, l__} :> {f, l}] & /@ Union[Flatten[data[[All, {1, 2}]]]]

(with the caveat tested on a hampster-powered loungebook...)

What will be fastest really depends on data structure (is it in fact in canonical subset order?), how it will be used (e.g. will you be only interested in some small subset of points, or do you need all distance metrics, etc.), so additional detail might flesh out a better method.

Edit: I'll go ahead and add a version that assumes the subset structure alluded to (it seems from your example and other statements this may be the case, if not, comment and I'll remove). It generates the same results as the Case solutions:

dist[data_] := Module[{dists = data[[All, 3]],
   pv, r, u = Range@data[[-1, -2]], n},
  n = Last@u;
  pv = Prepend[Accumulate[Range[n - 1, 1, -1]], 0];
  r = Range@n;

  Transpose[{Delete[r, #], 
      Join[dists[[pv[[;; # - 1]] + Range[# - 1, 1, -1]]], 
       dists[[Range[pv[[#]] + 1, pv[[#]] + (n - #)]]]]}] & /@ u];

Using

data = Subsets[Range@n, {2}];
data = Flatten /@ Transpose[{data, Range@Length@data}];

to generate dummy data in that form and comparing to the Cases solutions for n of 50 to 500 in steps of 50:

So, by n=500, this is ~400X faster (and ~30X faster than the Nearest generalization) - I stopped testing there for lack of cigars to smoke while waiting and fear of the loungebook bursting into flames... do note that on a "real" machine I'd expect 10-20X or more speedup of both timings.

I think there's some more optimization in this quickly cobbled-up realization, I'll ponder that.

I note you state "... the routine will be used heavily in lieu of storing the lists output." To me this implies getting all the results in one go is of lesser interest than the outright speed getting single-shots. For that, a sort of custom Nearest is probably the best bet:

distB[data_] := Module[{dists = data[[All, 3]],
    pv, r, u = Range@data[[-1, -2]], n},
   n = Last@u;
   pv = Prepend[Accumulate[Range[n - 1, 1, -1]], 0];
   r = Range@n;

   Function[{arg}, 
    Transpose[{Delete[r, arg], 
      Join[dists[[pv[[;; arg - 1]] + Range[arg - 1, 1, -1]]], 
       dists[[Range[pv[[arg]] + 1, pv[[arg]] + (n - arg)]]]]}]]];

Usage is similar to Nearest - assuming the data is in data,

myfn=distB[data];

creates a customized function for that dataset. Usage is then just:

myfn@target

where target is the desired point element. You only create the function once for a given dataset, then use it against your desired targets.

This is quite a bit faster than any of the above (about 2-3X faster to get all results in my limited testing against my fastest so far), and is below clock resolution even on the loungebook for single-shot queries on n=1024. It is ~1600X faster for single-shot queries for n=1024 vs the Cases solutions (loungebook timing caveats as always).

I'll update the bmark later - it's nearing sunrise here...

Answer 2

If I understand your question right, I'd use Cases:

Cases[data, {f___, 2, l__} :> {f, l}]

{{1, Subscript[a, 1]}, {3, Subscript[a, 5]}, {4, Subscript[a, 6]}, {5, Subscript[a, 7]}}

And to wrap it all up,

Cases[data, {f___, #, l__} :> {f, l}] & /@ 
 Union[Flatten[data[[All, {1, 2}]]]]

{{{2, Subscript[a, 1]}, {3, Subscript[a, 2]}, {4, Subscript[a, 3]}, {5, Subscript[a, 4]}}, {{1, Subscript[a, 1]}, {3, Subscript[a, 5]}, {4, Subscript[a, 6]}, {5, Subscript[a, 7]}}, {{1, Subscript[ a, 2]}, {2, Subscript[a, 5]}, {4, Subscript[a, 8]}, {5, Subscript[ a, 9]}}, {{1, Subscript[a, 3]}, {2, Subscript[a, 6]}, {3, Subscript[a, 8]}, {5, Subscript[a, 10]}}, {{1, Subscript[a, 4]}, {2, Subscript[a, 7]}, {3, Subscript[a, 9]}, {4, Subscript[a, 10]}}}

You can sort by using SortBy.

Answer 3

I think that a better performance can be achieved by simplifying the matching pattern as below:

Join[
    Cases[data, {2, f_, l_} :> {f, l}], 
    Cases[data, {f_, 2, l_} :> {f, l}]
]

On my pc this is 25% faster.