mol-data: OrderedSet.indexedIntersect

src/mol-data/int/_spec/ordered-set.spec.ts

@@ -6,6 +6,7 @@
 
 import OrderedSet from '../ordered-set'
 import Interval from '../interval'
+import SortedArray from '../sorted-array';
 
 describe('ordered set', () => {
     function ordSetToArray(set: OrderedSet) {
@@ -163,6 +164,11 @@ describe('ordered set', () => {
     testEq('intersect AA', OrderedSet.intersect(arr136, OrderedSet.ofSortedArray([2, 3, 4, 6, 7])), [3, 6]);
     it('intersect AA1', () => expect(OrderedSet.union(arr136, OrderedSet.ofSortedArray([1, 3, 6]))).toBe(arr136));
 
+    testEq('idxIntersect 1', OrderedSet.indexedIntersect(
+        OrderedSet.ofSortedArray([1, 2, 4]),
+        SortedArray.ofSortedArray([1, 2, 3, 4, 5, 6]),
+        SortedArray.ofSortedArray([2, 4, 5, 8])), [0, 2]);
+
     testEq('subtract ES', OrderedSet.subtract(empty, singleton10), []);
     testEq('subtract ER', OrderedSet.subtract(empty, range1_4), []);
     testEq('subtract EA', OrderedSet.subtract(empty, arr136), []);

src/mol-data/int/impl/ordered-set.ts

@@ -286,4 +286,48 @@ export function forEach(set: OrderedSetImpl, f: (value: number, i: number, ctx:
         }
     }
     return ctx;
+}
+
+
+export function indexedIntersect(idxA: OrderedSetImpl, a: S, b: S): OrderedSetImpl {
+    if (a === b) return idxA;
+    const lenI  = size(idxA), lenA = a.length, lenB = b.length;
+    if (lenI === 0 || lenA === 0 || lenB === 0) return Empty;
+
+    // const { startI, startJ, endI, endJ } = getSuitableIntersectionRange(a, b);
+    // const commonCount = getCommonCount(a, b, startI, startJ, endI, endJ);
+
+    const startJ = S.findPredecessorIndex(b, a[min(idxA)]);
+    const endJ = S.findPredecessorIndex(b, a[max(idxA)] + 1);
+
+    let commonCount = 0;
+
+    let offset = 0;
+    let O = 0;
+    let j = startJ;
+    while (O < lenI && j < endJ) {
+        const x = a[getAt(idxA, O)], y = b[j];
+        if (x < y) { O++; }
+        else if (x > y) { j++; }
+        else { commonCount++; O++; j++; }
+    }
+
+    // no common elements
+    if (commonCount === 0) return Empty;
+    // A === B
+    if (commonCount === lenA && commonCount === lenB) return idxA;
+
+    const indices = new Int32Array(commonCount);
+
+    offset = 0;
+    O = 0;
+    j = startJ;
+    while (O < lenI && j < endJ) {
+        const x = a[getAt(idxA, O)], y = b[j];
+        if (x < y) { O++; }
+        else if (x > y) { j++; }
+        else { indices[offset++] = j; O++; j++; }
+    }
+
+    return ofSortedArray(indices);
 }

src/mol-data/int/ordered-set.ts

@@ -39,6 +39,7 @@ namespace OrderedSet {
 
     export const union: <T extends number = number>(a: OrderedSet<T>, b: OrderedSet<T>) => OrderedSet<T> = Base.union as any;
     export const intersect: <T extends number = number>(a: OrderedSet<T>, b: OrderedSet<T>) => OrderedSet<T> = Base.intersect as any;
+    export const indexedIntersect: <T extends number = number, S extends number = number>(idxA: OrderedSet<T>, a: SortedArray<S>, b: SortedArray<S>) => OrderedSet<T> = Base.indexedIntersect as any;
     /** Returns elements of `a` that are not in `b`, i.e `a` - `b` */
     export const subtract: <T extends number = number>(a: OrderedSet<T>, b: OrderedSet<T>) => OrderedSet<T> = Base.subtract as any;