# Efficient extraction of a subset

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.

• So the natural solution would be Cases, Select, GatherBy... Are those not fast enough for your purpose?
– kale
Commented Jun 25, 2015 at 22:23
• Is the answer by kale you've tested fast enough (it can be bettered significantly)? If it is, consider up-voting it and accepting it if appropriate.
– ciao
Commented Jun 25, 2015 at 23:18
• @ciao Significantly??? :)
– kale
Commented Jun 25, 2015 at 23:34
• @kale: 5 to 6X faster at least...
– ciao
Commented Jun 25, 2015 at 23:43
• Is data[[All, {1, 2}]] == Subsets[Range@Max@data[[All, {1, 2}]], {2}] always True? If so you can bypass pattern matching entirely. Commented Jun 26, 2015 at 0:58

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...

• You didn't disappoint.
– kale
Commented Jun 26, 2015 at 12:14
• Kale's 10X improvement was impressive; these 10^k improvements bring tears to my eyes. Excellent presentation. Many thanks. Commented Jun 26, 2015 at 15:03
• @kale: Under-promise, over-deliver ;-}
– ciao
Commented Jun 27, 2015 at 6:54

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.

• Performs as needed. Comparing times on larger data set. Definitely an improvement. Thanks. Commented Jun 25, 2015 at 22:39
• For n = 1024: you = 0.32 s; me = 3.5 s. More than an order of magnitude improvement. Thanks x 10. Commented Jun 25, 2015 at 22:55

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.

• Interesting observation. Rerunning the data for n = 1024, kale = 0.31 s, user8074 = 0.30 s on my machine. (For testing purposes one could use random numbers for the distances.) Now I wonder about running two threads and combining them.... Commented Jun 25, 2015 at 23:35