7
$\begingroup$

Given a large list of $n$ dimensional sublists where each sublist has integer elements I am trying to find the quickest way of finding all sublists that satisfy some constraint on their elements.

A typical constraint would be that a sublist sums to zero.

I have tried three ways - Pick with Unitize, Pick with Clip and compiling Select.

They all seem to have similar timings. The two Pick methods are somewhat faster for $n<6$ but compiling Select is faster for larger $n$.

Does anyone have a faster way of doing this? Especially when $n \ge6$?

(* n = 3 case *)

arr = RandomInteger[{-100, 100}, {10000000, 3}];

constraint = #[[1]] + #[[2]] + #[[3]] == 0 &;
compile = 
  With[{criteria = constraint}, 
   Compile[{{input, _Integer, 2}}, Select[input, criteria], 
    CompilationTarget -> "C", RuntimeOptions -> "Speed"]];

try1 = Pick[arr, 
     Unitize[#[[All, 1]] + #[[All, 2]] + #[[All, 3]] &@arr], 0]; // 
  RepeatedTiming // First

try2 = compile[arr]; // RepeatedTiming // First

try3 = Pick[arr, 
     1 - 
      Clip[#[[All, 1]] + #[[All, 2]] + #[[All, 3]] &@arr, {0, 0}, {1, 
        1}], 1]; // RepeatedTiming // First

try1 == try2 == try3

(*
0.200528

0.379881

0.222603

True
*)

and

(* n = 9 case *)

arr2 = RandomInteger[{-100, 100}, {10000000, 9}];

constraint = #[[1]] + #[[2]] + #[[3]] + #[[4]] + #[[5]] + #[[6]] + \
#[[7]] + #[[8]] + #[[9]] == 0 &;
compile = 
  With[{criteria = constraint}, 
   Compile[{{input, _Integer, 2}}, Select[input, criteria], 
    CompilationTarget -> "C", RuntimeOptions -> "Speed"]];

try1 = Pick[arr2, 
     Unitize[#[[All, 1]] + #[[All, 2]] + #[[All, 3]] + #[[All, 
           4]] + #[[All, 5]] + #[[All, 6]] + #[[All, 7]] + #[[All, 
           8]] + #[[All, 9]] &@arr2], 0]; // RepeatedTiming // First

try2 = compile[arr2]; // RepeatedTiming // First

try3 = Pick[arr2, 
     1 - 
      Clip[#[[All, 1]] + #[[All, 2]] + #[[All, 3]] + #[[All, 
            4]] + #[[All, 5]] + #[[All, 6]] + #[[All, 7]] + #[[All, 
            8]] + #[[All, 9]] &@arr2, {0, 0}, {1, 1}], 1]; // 
  RepeatedTiming // First

try1 == try2 == try3

(*
0.59877

0.446898

0.615144

True
*)

I am not sure of the etiquette on here in regard to extending the question and have received some great responses in the comments so far.

But as a matter of interest - instead of the summing to zero constraint - what would be a fast method of constraining the first $n-1$ elements to be in ascending order?

i.e $ element[[1]] \le element[[2]] \le \dots element[[n-1]]$ in a sublist?

$\endgroup$
6
  • 7
    $\begingroup$ you can also try Pick[arr, Total[arr, {2}], 0] $\endgroup$
    – kglr
    Jan 6, 2022 at 11:41
  • 4
    $\begingroup$ Pick[arr, arr.Table[1, Length[arr[[1]]]], 0] $\endgroup$
    – chyanog
    Jan 6, 2022 at 11:44
  • $\begingroup$ Before I anoint you both as code Jedis - can you give some insight as to why your methods are so much faster? I thought I was being moderately clever (well not clever in regard to just compiling Select LOL). $\endgroup$
    – 1729taxi
    Jan 6, 2022 at 11:52
  • 3
    $\begingroup$ @chyanog you could cut that down further and remove the Length by doing arr . (1 & /@ arr[[1]]) . Also might be worth checking ConstantArray vs Table, e . g arr.ConstantArray[1, n] $\endgroup$
    – flinty
    Jan 6, 2022 at 12:00
  • 2
    $\begingroup$ Total was developed by Anton Antonov under the supervision of Rob Knapp. See here for an interesting post by AA. Before the development of Total, a trick was to use Tr (now considered slow). See the entry for Tr in Clever little programs by Ted Ersek. $\endgroup$
    – user1066
    Jan 6, 2022 at 16:08

1 Answer 1

6
$\begingroup$

Here is another implementation for selecting rows that sum to 0. It is designed to parallelize well and to avoid Pick.

n = 9;
arr2 = RandomInteger[{-100, 100}, {10000000, n}];

threadCount = "ParallelThreadNumber" /. ("ParallelOptions" /. SystemOptions["ParallelOptions"]);
jobptr = 1 + Ceiling[(Length[arr2]) Subdivide[0., 1., threadCount]];

cf = Block[{input, i},
   With[{code = Sum[Compile`GetElement[input, i, j], {j, 1, n}] == 0},
    Compile[{{input, _Integer, 2}, {begin, _Integer}, {end, _Integer}},
     Block[{bag},
      bag = Internal`Bag[Most[{0}]];
      Do[If[code, Internal`StuffBag[bag, i]], {i, begin, end - 1}];
      Internal`BagPart[bag, All]
      ],
     CompilationTarget -> "C",
     RuntimeAttributes -> {Listable},
     Parallelization -> True,
     RuntimeOptions -> "Speed"]
    ]
   ];

result0 = Pick[arr2, Total[arr2, {2}], 0]; // RepeatedTiming // First
result1 = Pick[arr2, arr2.ConstantArray[1, n], 0]; // RepeatedTiming // First
result2 = arr2[[Join @@ cf[arr2, Most[jobptr], Rest[jobptr]]]]; // 
  RepeatedTiming // First

result0 == result1 == result2

0.279332

0.105336

0.0365174

True

And here are some ideas for the extended question (picking rows that are ordered up to the penultimate element):

cg = With[{m = n - 1},
   Compile[{{input, _Integer, 2}, {begin, _Integer}, {end, _Integer}},
    Block[{b, c, x, y, bag},
     bag = Internal`Bag[Most[{0}]];
     Do[
      b = True;
      c = 1;
      x = Compile`GetElement[input, i, 1];
      While[b && c < m,
       c++;
       y = Compile`GetElement[input, i, c];
       b = x <= y;
       x = y;
       ];
      If[b, Internal`StuffBag[bag, i]],
      {i, begin, end - 1}];
     Internal`BagPart[bag, All]
     ],
    CompilationTarget -> "C",
    RuntimeAttributes -> {Listable},
    Parallelization -> True,
    RuntimeOptions -> "Speed"
    ]
   ];

A = SparseArray[{Band[{1, 1}] -> -1., Band[{2, 1}] -> 1.}, {n, n - 2}];
one = ConstantArray[1, n - 2];


result0 = Select[arr2, OrderedQ[Most[#]] &]; // RepeatedTiming // First
result1 = Pick[arr2, UnitStep[arr2 . A] . one, n - 2]; // 
  RepeatedTiming // First
result2 = arr2[[Join @@ cg[arr2, Most[jobptr], Rest[jobptr]]]]; // 
  RepeatedTiming // First

result0 == result1 == result2

10.6702

0.812736

0.0381827

True

The approach for result1 is to use a SparseArray A to apply the differences in each row. Then we apply UnitStep and sum the resulting entries to count the number of nonnegative entries (which correspond to an monotonically increasing pair of entries in the original matrix).

For result2, we again use the parallelized CompiledFunction cg. The While loop terminates early if it finds a pair of consecutive entries that is not monotonically increasing. Internal`Bag (along with Internal`StuffBag and Internal`BagPart) is an undocumented container type that allows efficient appending of elements (it's probably just using a C++ std::vector and its push_back routine as backend).

$\endgroup$
2
  • $\begingroup$ What can I say but fantastic answer. This part of my coding was my bottleneck from a speed perspective - it no longer isn't. Sadly for me the bottleneck is now something Mathematica has no control over. But I have gone from requiring several hours to five minutes or so (though I have jobs that still take a day or so but were previously going to take weeks) - with your recommendations plus some things I was doing in other areas of my code. Thank you very much Henrik. $\endgroup$
    – 1729taxi
    Jan 10, 2022 at 14:17
  • $\begingroup$ I am pleased to hear that! Good luck for your project! $\endgroup$ Jan 10, 2022 at 16:30

Your Answer

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

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