allow rank to take a comparator

README.md

@@ -141,13 +141,15 @@ An optional *accessor* function may be specified, which is equivalent to calling
 
 Similar to *quantile*, but expects the input to be a **sorted** *array* of values. In contrast with *quantile*, the accessor is only called on the elements needed to compute the quantile.
 
-<a name="rank" href="#rank">#</a> d3.<b>rank</b>(<i>iterable</i>[, <i>accessor</i>]) · [Source](https://github.com/d3/d3-array/blob/main/src/rank.js), [Examples](https://observablehq.com/@d3/rank)
+<a name="rank" href="#rank">#</a> d3.<b>rank</b>(<i>iterable</i>[, <i>comparator</i>]) · [Source](https://github.com/d3/d3-array/blob/main/src/rank.js), [Examples](https://observablehq.com/@d3/rank)
+<br><a name="rank" href="#rank">#</a> d3.<b>rank</b>(<i>iterable</i>[, <i>accessor</i>])
 
-Returns an array with the rank of each value in the *iterable*, *i.e.* the index of the value when the iterable is sorted. Nullish values are sorted to the end and ranked NaN. An optional *accessor* function may be specified, which is equivalent to calling *array*.map(*accessor*) before computing the ranks. Ties (equivalent values) all get the same rank, defined as the first time the value is found.
+Returns an array with the rank of each value in the *iterable*, *i.e.* the zero-based index of the value when the iterable is sorted. Nullish values are sorted to the end and ranked NaN. An optional *comparator* or *accessor* function may be specified; the latter is equivalent to calling *array*.map(*accessor*) before computing the ranks. If *comparator* is not specified, it defaults to [ascending](#ascending). Ties (equivalent values) all get the same rank, defined as the first time the value is found.
 
 ```js
 d3.rank([{x: 1}, {}, {x: 2}, {x: 0}], d => d.x); // [1, NaN, 2, 0]
 d3.rank(["b", "c", "b", "a"]); // [1, 3, 1, 0]
+d3.rank([1, 2, 3], d3.descending); // [2, 1, 0]
 ```
 
 <a name="variance" href="#variance">#</a> d3.<b>variance</b>(<i>iterable</i>[, <i>accessor</i>]) · [Source](https://github.com/d3/d3-array/blob/main/src/variance.js), [Examples](https://observablehq.com/@d3/d3-mean-d3-median-and-friends)

src/rank.js

@@ -1,23 +1,24 @@
-import range from "./range.js";
-import sort from "./sort.js";
+import ascending from "./ascending.js";
+import {ascendingDefined, compareDefined} from "./sort.js";
 
-export default function rank(values, valueof) {
+export default function rank(values, valueof = ascending) {
   if (typeof values[Symbol.iterator] !== "function") throw new TypeError("values is not iterable");
-  values = Array.from(values, valueof);
-  const n = values.length;
-  const r = new Float64Array(n);
-  let last, l;
-  sort(range(n), (i) => values[i]).forEach((j, i) => {
-    const value = values[j];
-    if (value == null || !(value <= value)) {
-      r[j] = NaN;
-      return;
-    }
-    if (last === undefined || !(value <= last)) {
-      last = value;
-      l = i;
-    }
-    r[j] = l;
-  });
-  return r;
+  let V = Array.from(values);
+  const R = new Float64Array(V.length);
+  if (valueof.length === 1) V = V.map(valueof), valueof = ascending;
+  const compareIndex = (i, j) => valueof(V[i], V[j]);
+  let k, r;
+  Uint32Array
+    .from(V, (_, i) => i)
+    .sort(valueof === ascending ? (i, j) => ascendingDefined(V[i], V[j]) : compareDefined(compareIndex))
+    .forEach((j, i) => {
+      const c = compareIndex(j, k === undefined ? j : k);
+      if (c >= 0) {
+        if (k === undefined || c > 0) k = j, r = i;
+        R[j] = r;
+      } else {
+        R[j] = NaN;
+      }
+    });
+  return R;
 }

test/rank-test.js

@@ -1,4 +1,6 @@
 import assert from "assert";
+import ascending from "../src/ascending.js";
+import descending from "../src/descending.js";
 import rank from "../src/rank.js";
 
 it("rank(numbers) returns the rank of numbers", () => {
@@ -28,12 +30,19 @@ it("rank(values, valueof) accepts an accessor", () => {
   assert.deepStrictEqual(rank([{x: 3}, {x: 1}, {x: 2}, {x: 4}, {}], d => d.x), Float64Array.of(2, 0, 1, 3, NaN));
 });
 
-it("rank(values, ties) computes the ties as expected", () => {
+it("rank(values, compare) accepts a comparator", () => {
+  assert.deepStrictEqual(rank([{x: 3}, {x: 1}, {x: 2}, {x: 4}, {}], (a, b) => a.x - b.x), Float64Array.of(2, 0, 1, 3, NaN));
+  assert.deepStrictEqual(rank([{x: 3}, {x: 1}, {x: 2}, {x: 4}, {}], (a, b) => b.x - a.x), Float64Array.of(1, 3, 2, 0, NaN));
+  assert.deepStrictEqual(rank(["aa", "ba", "bc", "bb", "ca"], (a, b) => ascending(a[0], b[0]) || ascending(a[1], b[1])), Float64Array.of(0, 1, 3, 2, 4));
+  assert.deepStrictEqual(rank(["A", null, "B", "C", "D"], descending), Float64Array.of(3, NaN, 2, 1, 0));
+});
+
+it("rank(values) computes the ties as expected", () => {
   assert.deepStrictEqual(rank(["a", "b", "b", "b", "c"]), Float64Array.of(0, 1, 1, 1, 4));
   assert.deepStrictEqual(rank(["a", "b", "b", "b", "b", "c"]), Float64Array.of(0, 1, 1, 1, 1, 5));
 });
 
-it("rank(values, ties) handles NaNs as expected", () => {
+it("rank(values) handles NaNs as expected", () => {
   assert.deepStrictEqual(rank(["a", "b", "b", "b", "c", null]), Float64Array.of(0, 1, 1, 1, 4, NaN));
   assert.deepStrictEqual(rank(["a", "b", "b", "b", "b", "c", null]), Float64Array.of(0, 1, 1, 1, 1, 5, NaN));
 });