makeBuckets

src/helpers.d.ts

@@ -14,4 +14,6 @@ declare module Helpers {
     export type UintArray = Uint8Array | Uint16Array | Uint32Array | number[]
     export type ValueOf<T> = T[keyof T]
     export type ArrayCtor<T> = { new(size: number): { [i: number]: T, length: number } }
+    /** assignable ArrayLike version */
+    export type ArrayLike<T> =  { [i: number]: T, length: number }
 }

src/mol-data/util.ts

@@ -5,6 +5,7 @@
  */
 
 export * from './util/chunked-array'
+export * from './util/buckets'
 export * from './util/equivalence-classes'
 export * from './util/hash-functions'
 export * from './util/sort'

src/mol-data/util/_spec/buckets.spec.ts

@@ -0,0 +1,35 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { createRangeArray } from '../array';
+import { makeBuckets } from '../buckets';
+
+describe('buckets', () => {
+
+    function reorder(order: ArrayLike<number>, data: any[]): any[] {
+        const ret = [];
+        for (const i of (order as number[])) ret[ret.length] = data[i];
+        return ret;
+    }
+
+    it('full range', () => {
+        const xs = [1, 1, 2, 2, 3, 1];
+        const range = createRangeArray(0, xs.length - 1);
+        const bs = makeBuckets(range, i => xs[i]);
+
+        expect(reorder(range, xs)).toEqual([1, 1, 1, 2, 2, 3]);
+        expect(Array.from(bs)).toEqual([0, 3, 5, 6]);
+    });
+
+    it('subrange', () => {
+        const xs = [2, 1, 2, 1, 2, 3, 1];
+        const range = createRangeArray(0, xs.length - 1);
+        const bs = makeBuckets(range, i => xs[i], 1, 5);
+
+        expect(reorder(range, xs)).toEqual([2, 1, 1, 2, 2, 3, 1]);
+        expect(Array.from(bs)).toEqual([1, 3, 5]);
+    })
+})

src/mol-data/util/buckets.ts

@@ -0,0 +1,74 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+type Bucket = {
+    count: number,
+    offset: number
+}
+
+function _makeBuckets(indices: Helpers.ArrayLike<number>, getKey: (i: number) => any, start: number, end: number) {
+    const buckets = new Map<any, Bucket>();
+    const bucketList: Bucket[] = [];
+
+    let prevKey = getKey(indices[0]);
+    let isBucketed = true;
+    for (let i = start; i < end; i++) {
+        const key = getKey(indices[i]);
+        if (buckets.has(key)) {
+            buckets.get(key)!.count++;
+            if (prevKey !== key) isBucketed = false;
+        } else {
+            const bucket: Bucket = { count: 1, offset: i };
+            buckets.set(key, bucket);
+            bucketList[bucketList.length] = bucket;
+        }
+        prevKey = key;
+    }
+
+    const bucketOffsets = new Int32Array(bucketList.length + 1);
+    bucketOffsets[bucketList.length] = end;
+
+    if (isBucketed) {
+        for (let i = 0; i < bucketList.length; i++) bucketOffsets[i] = bucketList[i].offset;
+        return bucketOffsets;
+    }
+
+    let offset = 0;
+    for (let i = 0; i < bucketList.length; i++) {
+        const b = bucketList[i];
+        b.offset = offset;
+        offset += b.count;
+    }
+
+    const reorderedIndices = new Int32Array(end - start);
+    for (let i = start; i < end; i++) {
+        const key = getKey(indices[i]);
+        const bucket = buckets.get(key)!;
+        reorderedIndices[bucket.offset++] = indices[i];
+    }
+
+    for (let i = 0, _i = reorderedIndices.length; i < _i; i++) {
+        indices[i + start] = reorderedIndices[i];
+    }
+
+    bucketOffsets[0] = start;
+    for (let i = 1; i < bucketList.length; i++) bucketOffsets[i] = bucketList[i - 1].offset + start;
+
+    return bucketOffsets;
+}
+
+/**
+ * Reorders indices so that the same keys are next to each other, [start, end)
+ * Returns the offsets of buckets. So that [offsets[i], offsets[i + 1]) determines the range.
+ */
+export function makeBuckets<T>(indices: Helpers.ArrayLike<number>, getKey: (i: number) => string | number, start?: number, end?: number): ArrayLike<number> {
+    const s = start || 0;
+    const e = typeof end === 'undefined' ? indices.length : end;
+
+    if (e - s <= 0) throw new Error('Can only bucket non-empty collections.');
+
+    return _makeBuckets(indices, getKey, s, e);
+}