Which elements of a large array are part of a set?

Question

I have a large integer packed array arr. My array is multidimensional, but for simplicity let's consider 1D arrays for now. ArrayReshape takes care of the rest.

I also have a set of integers (a list).

The task is to replace those elements of arr which are in set with 1, and those which aren't with 0.

What is the best way to do it, where "best" means fastest for as long we don't run out of memory?

For testing purposes, let's use

n = 100000000;
k = 20;
arr = RandomInteger[100, n];
set = RandomInteger[100, k];

---

What have I tried so far?

Method 1. For a single element of set, denoted elem, we can test using 1 - Unitize[arr - elem]. We can test for each element of set one by one using

res = 1 - Fold[# Unitize[arr - #2] &, ConstantArray[1, Length[arr]], set];

This is linear in the size of set, and we can do better than that. I used Fold instead of mapping over set to avoid having to simultaneously store a huge array in memory for each element of set.

Method 2. An alternative is using associations for looking up the set elements, which should be no worse than logarithmic complexity in the size of set (i.e. much better than linear).

ass = AssociationThread[set -> ConstantArray[1, Length[set]]]
res = Lookup[ass, arr, 0];

It's clear that for large enough k this will be faster than method 1 due to its better complexity. In practice the threshold where that happens is only $k=5$ for my test data. But this method does not produce a packed array result, which is a disadvantage!

For comparison, I also implemented the same thing in C++ (using a naive lookup with std::set). For $k=5$ I get the following timings: 5.9 s for method 1; 5.9 s for method 2 (re-packing the array takes an additional 0.7 s); 1.4 s for C++.

Answer 1

Here is a version that seems pretty fast and does not involve compiled code:

ClearAll[getMask];
getMask[arr_, set_] :=
  Module[{max = Max[arr], min = Min[arr], inds, dim, lset},
    lset = Pick[set, UnitStep[set - min]*UnitStep[max - set], 1];
    dim = max - min + 1;
    inds = SparseArray[{}, dim];
    inds[[lset - min + 1]] = 1;
    inds[[arr - min + 1]]
  ]

It uses the main idea similar to the one I used in my answer to another question of yours. Basically, we create an index array, where your original data serve as indices.

It returns the result on a test data of the same sizes as in your question, in about 3 sec. on my machine. An added advantage is that the data is returned already in the SparseArray form, which in some cases may be a lot more memory-efficient.

In now deleted answer, WReach used a similar method based on simple integer arrays. I didn't use it here intentionally, due to memory considerations, although it is slightly faster. However, it just occurred to me that since here there are no indications that the data is sparse, one can use that too, without much memory overhead, in general. You can look at my answer I linked above, where I used both and introduced automatic method switch based on the sparseness parameter. Something very similar can be done here, if needed.

Also, while I didn't do it, but one can optimize this further in cases when all integers are strictly positive.

Answer 2

I tried to find a solution that uses the vectorization capabilities of Compiled, because in its core, this problem is very nice parallelizable. Unfortunately, I got only very inconsistent timings and it seems to be hard to know upfront why a certain implementation is faster or slower.

Let's assume we want to compute one element i from the array arr, then there are several ways to do this. One way is

1 - Min[Unitize[set - i]]

This code works similar to what Szabolcs is doing, but using Min instead of a multiplication seemed a bit faster. Another way to compute one element is a very simple loop that goes through all elements of set and returns 1 when it has found i in set.

Module[{res = 0},
 Do[
  If[i == j,
   res = 1;
   Break[];
   ], {j, set}
  ];
 res
 ]

Let's use these two bodies inside a compiled function:

fc1 = Compile[{{i, _Integer, 0}, {set, _Integer, 1}},
   Module[{res = 0},
    Do[
     If[i == j,
      res = 1;
      Break[];
      ], {j, set}
     ];
    res
    ],
   Parallelization -> True,
   RuntimeAttributes -> {Listable},
   RuntimeOptions -> "Speed",
   CompilationTarget -> "C"
   ];

fc2 = Compile[{{i, _Integer, 0}, {set, _Integer, 1}},
   1 - Min[Unitize[set - i]],
   Parallelization -> True,
   RuntimeAttributes -> {Listable},
   RuntimeOptions -> "Speed"
   ];

Measuring the runtime against Szabolcs method using a 3d arr instead of a one dimensional list:

n = 500;
k = 20;
arr = RandomInteger[100, {n, n, n}];
set = RandomInteger[100, k];

res = 1 - Fold[# Unitize[arr - #2] &, 
  ConstantArray[1, Length[arr]], set]; // AbsoluteTiming

The calculation of res took 5.06 s on my machine.

res2 = fc1[arr, set]; // AbsoluteTiming
res3 = fc2[arr, set]; // AbsoluteTiming

While the calculation of res2 takes 5.46 s, res3 only uses 4.45 s which is over half a second faster.

Before you ask, yes it was intentional to leave the CompilationTarget in fc2 out. It's just faster that way on my machine. fc1 on the other hand becomes even slower if it is not compiled to "C".

That's not the end of the story! If I change arr now and use a simple one dimensional array of size 10^8 like Szabolcs did, then the timing is really weird.

Szabolcs' method needs 4.06 s, fc1 now only needs 3.21 s and fc2 needs a runtime of 4.64 s.

Maybe you like to try this yourself.