diff --git a/src/helpers.d.ts b/src/helpers.d.ts index 69dca25924d4e2fc6dd927672674506e1ad935ff..9ba96996e0f9eacc9c0ce3a3baa16f7861670723 100644 --- a/src/helpers.d.ts +++ b/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 } } \ No newline at end of file diff --git a/src/mol-data/util.ts b/src/mol-data/util.ts index c541a1e9e9dc71678081bc32c4f158fda3b18b39..e92041312001c8f55e84f79b3f49dd5d9c16e109 100644 --- a/src/mol-data/util.ts +++ b/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' diff --git a/src/mol-data/util/_spec/buckets.spec.ts b/src/mol-data/util/_spec/buckets.spec.ts new file mode 100644 index 0000000000000000000000000000000000000000..2207a53f937fc7e5a17691bee96e08c26dc3825b --- /dev/null +++ b/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]); + }) +}) \ No newline at end of file diff --git a/src/mol-data/util/buckets.ts b/src/mol-data/util/buckets.ts new file mode 100644 index 0000000000000000000000000000000000000000..58a747e9dda5863bb5c41093e7d71a15f6235df5 --- /dev/null +++ b/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); +} \ No newline at end of file