Pick with Unitize vs Pick with Clip vs Compiled Select speed - anything faster?

Question

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?

Answer 1

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