Top-k and selection

Often you do not need everything sorted: only the ten best scores, the median, or the slowest requests. These calls do much less work than a full sort.

You want	Use
The `k` largest or smallest values, and where they were	`top_k`
Only the positions of the `k` largest or smallest	`argselect`
The value at one position of the sorted order (median, percentile)	`select`
The first `k` positions sorted, the rest in any order	`partial_sort`

The k largest or smallest values

top_k(a, k) returns two arrays: the k smallest values in order, and their positions in a. Ask for a descending order to get the k largest. a is not changed.

import numpy as np
import kwker

response_ms = np.array([120, 85, 430, 95, 610, 150, 520])
slowest, where = kwker.top_k(response_ms, 3, descending=True)
print(slowest)
print(where)
fastest, where = kwker.top_k(response_ms, 2)
print(fastest, where)

[610 520 430]
[4 6 2]
[85 95] [1 3]

RustAdd kwker to Cargo.toml, then cargo run.

use kwker::Order;

fn main() {
    let response_ms = [120, 85, 430, 95, 610, 150, 520];
    let (slowest, at): (Vec<i32>, Vec<usize>) = kwker::top_k(&response_ms, 3, Order::DESCENDING, true);
    println!("{slowest:?}");
    println!("{at:?}");
    let (fastest, at): (Vec<i32>, Vec<usize>) = kwker::top_k(&response_ms, 2, Order::ASCENDING, true);
    println!("{fastest:?} {at:?}");
}

Output

[610, 520, 430]
[4, 6, 2]
[85, 95] [1, 3]

CInclude kwker.h; link with -lkwker_c.

#include <inttypes.h>
#include <stdio.h>
#include <kwker.h>

int main(void) {
    const int64_t response_ms[] = {120, 85, 430, 95, 610, 150, 520};
    int64_t slowest[3];
    uint64_t at[3];
    kwker_i64_top_k(response_ms, 7, 3, KWKER_DESCENDING, 1, slowest, at);
    printf("%" PRId64 " %" PRId64 " %" PRId64 "\n", slowest[0], slowest[1], slowest[2]);
    printf("%" PRIu64 " %" PRIu64 " %" PRIu64 "\n", at[0], at[1], at[2]);
    return 0;
}

Output

610 520 430
4 6 2

C++Include kwker.hpp; link with -lkwker_c.

#include <iostream>
#include <vector>
#include <kwker.hpp>

int main() {
    std::vector<int> response_ms{120, 85, 430, 95, 610, 150, 520};
    auto [slowest, at] = kwker::top_k(response_ms, 3, kwker::Order::descending);
    for (int v : slowest) std::cout << v << ' ';
    std::cout << '\n';
    for (auto i : at) std::cout << i << ' ';
    std::cout << '\n';
}

Output

610 520 430
4 6 2

const kwk = require("kwker");

const responseMs = new Int32Array([120, 85, 430, 95, 610, 150, 520]);
const slowest = kwk.topK(responseMs, 3, { descending: true });
console.log(slowest.values, slowest.indices);
const fastest = kwk.topK(responseMs, 2);
console.log(fastest.values, fastest.indices);

Int32Array(3) [ 610, 520, 430 ] Uint32Array(3) [ 4, 6, 2 ]
Int32Array(2) [ 85, 95 ] Uint32Array(2) [ 1, 3 ]

Gogo get kwker.io/go, then go run .

package main

import (
	"fmt"

	"kwker.io/go/kwker"
)

func main() {
	responseMs := []int32{120, 85, 430, 95, 610, 150, 520}
	slowest, where := kwker.TopK(responseMs, 3, kwker.Descending, true)
	fmt.Println(slowest)
	fmt.Println(where)
	fastest, where := kwker.TopK(responseMs, 2, kwker.Ascending, true)
	fmt.Println(fastest, where)
}

Output

[610 520 430]
[4 6 2]
[85 95] [1 3]

JavaPut kwker.jar on the class path.

import io.kwker.Kwker;
import java.util.Arrays;

public class Example {
    public static void main(String[] args) {
        int[] responseMs = {120, 85, 430, 95, 610, 150, 520};
        int[] slowest = Kwker.topK(responseMs, 3, Kwker.DESCENDING);   // positions
        System.out.println(Arrays.toString(Arrays.stream(slowest).map(i -> responseMs[i]).toArray()));
        System.out.println(Arrays.toString(slowest));
        int[] fastest = Kwker.topK(responseMs, 2, Kwker.ASCENDING);
        System.out.println(Arrays.toString(Arrays.stream(fastest).map(i -> responseMs[i]).toArray()) + " " + Arrays.toString(fastest));
    }
}

Output

[610, 520, 430]
[4, 6, 2]
[85, 95] [1, 3]

C#Add the Kwker NuGet package, then dotnet run.

using Kwker;

var responseMs = new int[] { 120, 85, 430, 95, 610, 150, 520 };
var (slowest, where) = Sorter.TopK<int>(responseMs, 3, Order.Descending);
Console.WriteLine(string.Join(" ", slowest));
Console.WriteLine(string.Join(" ", where));
var (fastest, at) = Sorter.TopK<int>(responseMs, 2);
Console.WriteLine(string.Join(" ", fastest) + " | " + string.Join(" ", at));

Output

610 520 430
4 6 2
85 95 | 1 3

import numpy as np
import kwker

points = np.array([10, 30, 30, 20, 30])
values, where = kwker.top_k(points, 2, descending=True)
print(values, where)

[30 30] [1 2]

RustAdd kwker to Cargo.toml, then cargo run.

use kwker::Order;

fn main() {
    let points = [10, 30, 30, 20, 30];
    let (values, at): (Vec<i32>, Vec<usize>) = kwker::top_k(&points, 2, Order::DESCENDING, true);
    println!("{values:?} {at:?}");
}

Output

[30, 30] [1, 2]

CInclude kwker.h; link with -lkwker_c.

#include <inttypes.h>
#include <stdio.h>
#include <kwker.h>

int main(void) {
    const int32_t points[] = {10, 30, 30, 20, 30};
    int32_t values[2];
    uint64_t at[2];
    kwker_i32_top_k(points, 5, 2, KWKER_DESCENDING, 1, values, at);
    printf("%d %d | %" PRIu64 " %" PRIu64 "\n", values[0], values[1], at[0], at[1]);
    return 0;
}

Output

30 30 | 1 2

C++Include kwker.hpp; link with -lkwker_c.

#include <iostream>
#include <vector>
#include <kwker.hpp>

int main() {
    std::vector<int> points{10, 30, 30, 20, 30};
    auto [values, at] = kwker::top_k(points, 2, kwker::Order::descending);
    std::cout << values[0] << ' ' << values[1] << " | " << at[0] << ' ' << at[1] << '\n';
}

Output

30 30 | 1 2

const kwk = require("kwker");

const { values, indices } = kwk.topK(new Int32Array([10, 30, 30, 20, 30]), 2, { descending: true });
console.log(values, indices);

Int32Array(2) [ 30, 30 ] Uint32Array(2) [ 1, 2 ]

Gogo get kwker.io/go, then go run .

package main

import (
	"fmt"

	"kwker.io/go/kwker"
)

func main() {
	points := []int32{10, 30, 30, 20, 30}
	values, where := kwker.TopK(points, 2, kwker.Descending, true)
	fmt.Println(values, where)
}

Output

[30 30] [1 2]

JavaPut kwker.jar on the class path.

import io.kwker.Kwker;
import java.util.Arrays;

public class Example {
    public static void main(String[] args) {
        int[] points = {10, 30, 30, 20, 30};
        System.out.println(Arrays.toString(Kwker.topK(points, 2, Kwker.DESCENDING)));
    }
}

Output

[1, 2]

C#Add the Kwker NuGet package, then dotnet run.

using Kwker;

var points = new int[] { 10, 30, 30, 20, 30 };
var (values, where) = Sorter.TopK<int>(points, 2, Order.Descending);
Console.WriteLine(string.Join(" ", values) + " | " + string.Join(" ", where));

Output

30 30 | 1 2

If you do not need the k values in order, ask for them unsorted (sorted=False in Python, false as the last argument in Rust, 0 in C). It can save a little time on large k.

Only the positions

argselect(a, k) returns the positions of the k smallest values (or largest, in descending order), in no particular order.

import numpy as np
import kwker

prices = np.array([9.5, 2.0, 7.25, 1.5, 8.0])
cheapest = kwker.argselect(prices, 2)
print(np.sort(cheapest))
print(prices[np.sort(cheapest)])

[1 3]
[2.  1.5]

RustAdd kwker to Cargo.toml, then cargo run.

use kwker::Order;

fn main() {
    let prices = [9.5, 2.0, 7.25, 1.5, 8.0];
    let mut cheapest: Vec<usize> = kwker::argselect(&prices, 2, Order::ASCENDING);
    cheapest.sort();
    let chosen: Vec<f64> = cheapest.iter().map(|&i| prices[i]).collect();
    println!("{cheapest:?}");
    println!("{chosen:?}");
}

Output

[1, 3]
[2.0, 1.5]

CInclude kwker.h; link with -lkwker_c.

#include <inttypes.h>
#include <stdio.h>
#include <kwker.h>

int main(void) {
    const double prices[] = {9.5, 2.0, 7.25, 1.5, 8.0};
    uint64_t cheapest[2];
    kwker_f64_argselect(prices, 5, 2, KWKER_ASCENDING, cheapest);
    kwker_u64_sort(cheapest, 2);
    printf("%" PRIu64 " %" PRIu64 "\n", cheapest[0], cheapest[1]);
    printf("%g %g\n", prices[cheapest[0]], prices[cheapest[1]]);
    return 0;
}

Output

1 3
2 1.5

C++Include kwker.hpp; link with -lkwker_c.

#include <iostream>
#include <vector>
#include <kwker.hpp>

int main() {
    std::vector<double> prices{9.5, 2.0, 7.25, 1.5, 8.0};
    std::vector<uint64_t> cheapest = kwker::argselect(prices, 2);
    kwker::sort(cheapest);
    for (auto i : cheapest) std::cout << i << ' ';
    std::cout << '\n';
    for (auto i : cheapest) std::cout << prices[i] << ' ';
    std::cout << '\n';
}

Output

1 3
2 1.5

Gogo get kwker.io/go, then go run .

package main

import (
	"fmt"
	"slices"

	"kwker.io/go/kwker"
)

func main() {
	prices := []float64{9.5, 2.0, 7.25, 1.5, 8.0}
	cheapest := kwker.ArgSelect(prices, 2, kwker.Ascending)
	slices.Sort(cheapest)
	fmt.Println(cheapest)
	for _, i := range cheapest {
		fmt.Print(prices[i], " ")
	}
	fmt.Println()
}

Output

[1 3]
2 1.5

JavaPut kwker.jar on the class path.

import io.kwker.Kwker;
import java.util.Arrays;

public class Example {
    public static void main(String[] args) {
        double[] prices = {9.5, 2.0, 7.25, 1.5, 8.0};
        int[] cheapest = Kwker.argselect(prices, 2, Kwker.ASCENDING);
        Arrays.sort(cheapest);
        System.out.println(Arrays.toString(cheapest));
        System.out.println(Arrays.toString(Arrays.stream(cheapest).mapToDouble(i -> prices[i]).toArray()));
    }
}

Output

[1, 3]
[2.0, 1.5]

C#Add the Kwker NuGet package, then dotnet run.

using Kwker;

var prices = new double[] { 9.5, 2.0, 7.25, 1.5, 8.0 };
var cheapest = Sorter.ArgSelect<double>(prices, 2);
Array.Sort(cheapest);
Console.WriteLine(string.Join(" ", cheapest));
Console.WriteLine(string.Join(" ", cheapest.Select(i => prices[i])));

Output

1 3
2 1.5

The median and other positions

select(a, k) moves the value that belongs at position k (counting from 0) into place k. Every value before it is smaller or equal, every value after it is larger or equal, but neither side is sorted. This is how you get a median or a percentile quickly.

import numpy as np
import kwker

a = np.array([31, 4, 15, 92, 65, 35, 89, 79, 26])
k = len(a) // 2
kwker.select(a, k)
print("median:", a[k])
print("smaller side:", np.sort(a[:k]))
print("larger side: ", np.sort(a[k + 1:]))

median: 35
smaller side: [ 4 15 26 31]
larger side:  [65 79 89 92]

RustAdd kwker to Cargo.toml, then cargo run.

fn main() {
    let mut a = [31, 4, 15, 92, 65, 35, 89, 79, 26];
    let k = a.len() / 2;
    kwker::select_nth(&mut a, k);
    println!("median: {}", a[k]);
    let (smaller, rest) = a.split_at_mut(k);
    let larger = &mut rest[1..];
    kwker::sort(smaller);
    kwker::sort(larger);
    println!("smaller side: {smaller:?}");
    println!("larger side:  {larger:?}");
}

Output

median: 35
smaller side: [4, 15, 26, 31]
larger side:  [65, 79, 89, 92]

CInclude kwker.h; link with -lkwker_c.

#include <stdio.h>
#include <kwker.h>

int main(void) {
    int32_t a[] = {31, 4, 15, 92, 65, 35, 89, 79, 26};
    size_t n = 9, k = n / 2;
    kwker_i32_select(a, n, k, KWKER_ASCENDING);
    printf("median: %d\n", a[k]);
    kwker_i32_sort(a, k);
    kwker_i32_sort(a + k + 1, n - k - 1);
    printf("smaller side: %d %d %d %d\n", a[0], a[1], a[2], a[3]);
    printf("larger side:  %d %d %d %d\n", a[5], a[6], a[7], a[8]);
    return 0;
}

Output

median: 35
smaller side: 4 15 26 31
larger side:  65 79 89 92

C++Include kwker.hpp; link with -lkwker_c.

#include <iostream>
#include <vector>
#include <kwker.hpp>

int main() {
    std::vector<int> a{31, 4, 15, 92, 65, 35, 89, 79, 26};
    size_t k = a.size() / 2;
    kwker::select(a, k);
    std::cout << "median: " << a[k] << '\n';
    kwker::sort(a.data(), k);
    kwker::sort(a.data() + k + 1, a.size() - k - 1);
    std::cout << "smaller side:";
    for (size_t i = 0; i < k; i++) std::cout << ' ' << a[i];
    std::cout << "\nlarger side: ";
    for (size_t i = k + 1; i < a.size(); i++) std::cout << ' ' << a[i];
    std::cout << '\n';
}

Output

median: 35
smaller side: 4 15 26 31
larger side:  65 79 89 92

const kwk = require("kwker");

const a = new Int32Array([31, 4, 15, 92, 65, 35, 89, 79, 26]);
const k = Math.floor(a.length / 2);
kwk.select(a, k);
console.log("median:", a[k]);
console.log("smaller side:", kwk.sort(a.slice(0, k)));
console.log("larger side: ", kwk.sort(a.slice(k + 1)));

median: 35
smaller side: Int32Array(4) [ 4, 15, 26, 31 ]
larger side:  Int32Array(4) [ 65, 79, 89, 92 ]

Gogo get kwker.io/go, then go run .

package main

import (
	"fmt"
	"slices"

	"kwker.io/go/kwker"
)

func main() {
	a := []int32{31, 4, 15, 92, 65, 35, 89, 79, 26}
	k := len(a) / 2
	kwker.Select(a, k, kwker.Ascending)
	fmt.Println("median:", a[k])
	smaller, larger := slices.Clone(a[:k]), slices.Clone(a[k+1:])
	slices.Sort(smaller)
	slices.Sort(larger)
	fmt.Println("smaller side:", smaller)
	fmt.Println("larger side: ", larger)
}

Output

median: 35
smaller side: [4 15 26 31]
larger side:  [65 79 89 92]

JavaPut kwker.jar on the class path.

import io.kwker.Kwker;
import java.util.Arrays;

public class Example {
    public static void main(String[] args) {
        int[] a = {31, 4, 15, 92, 65, 35, 89, 79, 26};
        int k = a.length / 2;
        Kwker.select(a, k, Kwker.ASCENDING);
        System.out.println("median: " + a[k]);
        int[] smaller = Arrays.copyOfRange(a, 0, k), larger = Arrays.copyOfRange(a, k + 1, a.length);
        Arrays.sort(smaller);
        Arrays.sort(larger);
        System.out.println("smaller side: " + Arrays.toString(smaller));
        System.out.println("larger side:  " + Arrays.toString(larger));
    }
}

Output

median: 35
smaller side: [4, 15, 26, 31]
larger side:  [65, 79, 89, 92]

C#Add the Kwker NuGet package, then dotnet run.

using Kwker;

var a = new int[] { 31, 4, 15, 92, 65, 35, 89, 79, 26 };
int k = a.Length / 2;
Sorter.Select<int>(a, k);
Console.WriteLine($"median: {a[k]}");
Console.WriteLine("smaller side: " + string.Join(" ", a[..k].Order()));
Console.WriteLine("larger side:  " + string.Join(" ", a[(k + 1)..].Order()));

Output

median: 35
smaller side: 4 15 26 31
larger side:  65 79 89 92

For the 90th percentile of n values, select position int(0.9 * (n - 1)).

Sort only the beginning

partial_sort(a, k) sorts the first k positions and leaves the rest in any order. Use it when you show the first page of a long sorted list.

import numpy as np
import kwker

a = np.array([50, 10, 40, 20, 30, 60, 0])
kwker.partial_sort(a, 3)
print(a[:3])

[ 0 10 20]

RustAdd kwker to Cargo.toml, then cargo run.

fn main() {
    let mut a = [50, 10, 40, 20, 30, 60, 0];
    kwker::partial_sort(&mut a, 3);
    println!("{:?}", &a[..3]);
}

Output

[0, 10, 20]

CInclude kwker.h; link with -lkwker_c.

#include <stdio.h>
#include <kwker.h>

int main(void) {
    int32_t a[] = {50, 10, 40, 20, 30, 60, 0};
    kwker_i32_partial_sort(a, 7, 3, KWKER_ASCENDING);
    printf("%d %d %d\n", a[0], a[1], a[2]);
    return 0;
}

Output

0 10 20

C++Include kwker.hpp; link with -lkwker_c.

#include <iostream>
#include <vector>
#include <kwker.hpp>

int main() {
    std::vector<int> a{50, 10, 40, 20, 30, 60, 0};
    kwker::partial_sort(a, 3);
    std::cout << a[0] << ' ' << a[1] << ' ' << a[2] << '\n';
}

Output

0 10 20

const kwk = require("kwker");

const a = new Int32Array([50, 10, 40, 20, 30, 60, 0]);
kwk.partialSort(a, 3);
console.log(a.subarray(0, 3));

Int32Array(3) [ 0, 10, 20 ]

Gogo get kwker.io/go, then go run .

package main

import (
	"fmt"

	"kwker.io/go/kwker"
)

func main() {
	a := []int32{50, 10, 40, 20, 30, 60, 0}
	kwker.PartialSort(a, 3, kwker.Ascending)
	fmt.Println(a[:3])
}

Output

[0 10 20]

JavaPut kwker.jar on the class path.

import io.kwker.Kwker;
import java.util.Arrays;

public class Example {
    public static void main(String[] args) {
        int[] a = {50, 10, 40, 20, 30, 60, 0};
        Kwker.partialSort(a, 3, Kwker.ASCENDING);
        System.out.println(Arrays.toString(Arrays.copyOf(a, 3)));
    }
}

Output

[0, 10, 20]

C#Add the Kwker NuGet package, then dotnet run.

using Kwker;

var a = new int[] { 50, 10, 40, 20, 30, 60, 0 };
Sorter.PartialSort<int>(a, 3);
Console.WriteLine(string.Join(" ", a[..3]));

Output

0 10 20

Top-k with a filter

top_k_masked(a, mask, k) takes the top k among the entries where mask is true. The positions refer to the full array.

import numpy as np
import kwker

score = np.array([0.9, 0.4, 0.8, 0.95, 0.1])
active = np.array([True, True, True, False, True])
values, where = kwker.top_k_masked(score, active, 2, descending=True)
print(values, where)

[0.9 0.8] [0 2]

RustAdd kwker to Cargo.toml, then cargo run.

use kwker::Order;

fn main() {
    let score = [0.9, 0.4, 0.8, 0.95, 0.1];
    let active = [true, true, true, false, true];
    let (values, at) = kwker::top_k_masked(&score, &active, 2, Order::DESCENDING, true);
    println!("{values:?} {at:?}");
}

Output

[0.9, 0.8] [0, 2]

CInclude kwker.h; link with -lkwker_c.

#include <inttypes.h>
#include <stdio.h>
#include <kwker.h>

int main(void) {
    const double score[] = {0.9, 0.4, 0.8, 0.95, 0.1};
    const uint8_t active[] = {1, 1, 1, 0, 1};
    double values[2];
    uint64_t at[2];
    size_t count;
    kwker_f64_top_k_masked(score, 5, active, 0, 2, KWKER_DESCENDING, 1, values, at, &count);
    for (size_t i = 0; i < count; i++) printf("%g at %" PRIu64 "\n", values[i], at[i]);
    return 0;
}

Output

0.9 at 0
0.8 at 2

C++Include kwker.hpp; link with -lkwker_c.

#include <iostream>
#include <vector>
#include <kwker.hpp>

int main() {
    std::vector<double> score{0.9, 0.4, 0.8, 0.95, 0.1};
    std::vector<uint8_t> active{1, 1, 1, 0, 1};
    auto [values, at] = kwker::top_k_masked(score.data(), score.size(), active.data(), false, 2,
                                                kwker::Order::descending);
    for (size_t i = 0; i < values.size(); i++) std::cout << values[i] << " at " << at[i] << '\n';
}

Output

0.9 at 0
0.8 at 2

const kwk = require("kwker");

const score = new Float64Array([0.9, 0.4, 0.8, 0.95, 0.1]);
const active = new Uint8Array([1, 1, 1, 0, 1]);
const { values, indices } = kwk.topKMasked(score, active, 2, { descending: true });
console.log(values, indices);

Float64Array(2) [ 0.9, 0.8 ] Uint32Array(2) [ 0, 2 ]

Gogo get kwker.io/go, then go run .

package main

import (
	"fmt"

	"kwker.io/go/kwker"
)

func main() {
	score := []float64{0.9, 0.4, 0.8, 0.95, 0.1}
	active := []bool{true, true, true, false, true}
	values, at := kwker.TopKMasked(score, active, 2, kwker.Descending, true)
	fmt.Println(values, at)
}

Output

[0.9 0.8] [0 2]

JavaPut kwker.jar on the class path.

import io.kwker.Kwker;
import java.util.Arrays;

public class Example {
    public static void main(String[] args) {
        double[] score = {0.9, 0.4, 0.8, 0.95, 0.1};
        boolean[] active = {true, true, true, false, true};
        Kwker.TopK<double[]> top = Kwker.topKMasked(score, active, 2, Kwker.DESCENDING);
        System.out.println(Arrays.toString(top.values) + " " + Arrays.toString(top.indices));
    }
}

Output

[0.9, 0.8] [0, 2]

C#Add the Kwker NuGet package, then dotnet run.

using Kwker;

var score = new[] { 0.9, 0.4, 0.8, 0.95, 0.1 };
var active = new[] { true, true, true, false, true };
var (values, at) = Sorter.TopKMasked<double>(score, active, 2, Order.Descending);
Console.WriteLine(string.Join(" ", values) + " | " + string.Join(" ", at));

Output

0.9 0.8 | 0 2

Top-k per group

top_k_by_group(a, groups, k) finds the top k entries of every group. It is the SQL pattern ROW_NUMBER() OVER (PARTITION BY group ORDER BY value) <= k. Groups are integer labels. The call returns three arrays: the group labels in ascending order, offsets, and positions. Group i's positions are positions[offsets[i]:offsets[i + 1]].

import numpy as np
import kwker

sales = np.array([300, 120, 450, 80, 200, 610])
store = np.array([1, 2, 1, 2, 1, 2])
labels, offsets, positions = kwker.top_k_by_group(sales, store, 2, descending=True)
for i, label in enumerate(labels):
    rows = positions[offsets[i]:offsets[i + 1]]
    print("store", label, "best sales:", sales[rows])

store 1 best sales: [450 300]
store 2 best sales: [610 120]

RustAdd kwker to Cargo.toml, then cargo run.

use kwker::Order;

fn main() {
    let sales = [300, 120, 450, 80, 200, 610];
    let store: [i64; 6] = [1, 2, 1, 2, 1, 2];
    let (labels, offsets, positions) = kwker::top_k_by_group(&sales, &store, 2, Order::DESCENDING);
    for (i, label) in labels.iter().enumerate() {
        let best: Vec<i32> = positions[offsets[i]..offsets[i + 1]].iter().map(|&p| sales[p as usize]).collect();
        println!("store {label} best sales: {best:?}");
    }
}

Output

store 1 best sales: [450, 300]
store 2 best sales: [610, 120]

CInclude kwker.h; link with -lkwker_c.

#include <inttypes.h>
#include <stdio.h>
#include <kwker.h>

int main(void) {
    const int32_t sales[] = {300, 120, 450, 80, 200, 610};
    const int64_t store[] = {1, 2, 1, 2, 1, 2};
    int64_t labels[6];
    uint64_t offsets[7], positions[6];
    size_t ngroups;
    kwker_i32_top_k_by_group(sales, 6, store, 2, KWKER_DESCENDING, labels, offsets, positions, &ngroups);
    for (size_t g = 0; g < ngroups; g++) {
        printf("store %" PRId64 " best sales:", labels[g]);
        for (uint64_t j = offsets[g]; j < offsets[g + 1]; j++) printf(" %d", sales[positions[j]]);
        printf("\n");
    }
    return 0;
}

Output

store 1 best sales: 450 300
store 2 best sales: 610 120

C++Include kwker.hpp; link with -lkwker_c.

#include <iostream>
#include <vector>
#include <kwker.hpp>

int main() {
    std::vector<int> sales{300, 120, 450, 80, 200, 610};
    std::vector<int64_t> store{1, 2, 1, 2, 1, 2};
    auto best = kwker::top_k_by_group(sales.data(), sales.size(), store.data(), 2, kwker::Order::descending);
    for (size_t g = 0; g < best.labels.size(); g++) {
        std::cout << "store " << best.labels[g] << " best sales:";
        for (auto j = best.offsets[g]; j < best.offsets[g + 1]; j++) std::cout << ' ' << sales[best.indices[j]];
        std::cout << '\n';
    }
}

Output

store 1 best sales: 450 300
store 2 best sales: 610 120

const kwk = require("kwker");

const sales = new Int32Array([300, 120, 450, 80, 200, 610]);
const store = new Int32Array([1, 2, 1, 2, 1, 2]);
const { labels, offsets, indices } = kwk.topKByGroup(sales, store, 2, { descending: true });
labels.forEach((label, g) => {
  const rows = indices.subarray(offsets[g], offsets[g + 1]);
  console.log("store", label, "best sales:", Array.from(rows, (r) => sales[r]));
});

store 1 best sales: [ 450, 300 ]
store 2 best sales: [ 610, 120 ]

Gogo get kwker.io/go, then go run .

package main

import (
	"fmt"

	"kwker.io/go/kwker"
)

func main() {
	sales := []int32{300, 120, 450, 80, 200, 610}
	store := []int64{1, 2, 1, 2, 1, 2}
	labels, offsets, positions := kwker.TopKByGroup(sales, store, 2, kwker.Descending)
	for i, label := range labels {
		best := []int32{}
		for _, p := range positions[offsets[i]:offsets[i+1]] {
			best = append(best, sales[p])
		}
		fmt.Println("store", label, "best sales:", best)
	}
}

Output

store 1 best sales: [450 300]
store 2 best sales: [610 120]

JavaPut kwker.jar on the class path.

import io.kwker.Kwker;
import java.util.Arrays;

public class Example {
    public static void main(String[] args) {
        int[] sales = {300, 120, 450, 80, 200, 610};
        long[] store = {1, 2, 1, 2, 1, 2};
        Kwker.Groups g = Kwker.topKByGroup(sales, store, 2, Kwker.DESCENDING);
        for (int i = 0; i < g.labels.length; i++) {
            int[] best = Arrays.stream(g.indices, g.offsets[i], g.offsets[i + 1]).map(p -> sales[p]).toArray();
            System.out.println("store " + g.labels[i] + " best sales: " + Arrays.toString(best));
        }
    }
}

Output

store 1 best sales: [450, 300]
store 2 best sales: [610, 120]

C#Add the Kwker NuGet package, then dotnet run.

using Kwker;

var sales = new[] { 300, 120, 450, 80, 200, 610 };
var store = new long[] { 1, 2, 1, 2, 1, 2 };
var (labels, positions) = Sorter.TopKByGroup<int>(sales, store, 2, Order.Descending);
for (int i = 0; i < labels.Length; i++)
    Console.WriteLine($"store {labels[i]} best sales: {string.Join(" ", positions[i].Select(p => sales[p]))}");

Output

store 1 best sales: 450 300
store 2 best sales: 610 120

When the groups are runs of one array, such as the neighbors of each node in a graph stored as CSR (compressed sparse rows), top_k_segments(a, offsets, k) takes the k largest values of every segment: segment i is a[offsets[i]:offsets[i + 1]]. It returns the values and their positions inside each segment.

import numpy as np
import kwker

scores = np.array([0.9, 0.1, 0.5, 0.7, 0.3, 0.8], dtype=np.float32)
offsets = [0, 3, 4, 6]  # three segments: 0-2, 3, 4-5
values, positions = kwker.top_k_segments(scores, offsets, 2)
print(values)
print(positions)

[[0.9 0.5]
 [0.7 0. ]
 [0.8 0.3]]
[[ 0  2]
 [ 0 -1]
 [ 1  0]]

RustAdd kwker to Cargo.toml, then cargo run.

use kwker::Order;

fn main() {
    let scores = [0.9f32, 0.1, 0.5, 0.7, 0.3, 0.8];
    let offsets = [0usize, 3, 4, 6]; // three segments: 0-2, 3, 4-5
    let k = 2;
    let mut values = vec![0.0f32; 3 * k];
    let mut positions = vec![0usize; 3 * k];
    kwker::top_k_segments(&scores, &offsets, k, Order::DESCENDING, true, &mut values, &mut positions);
    for s in 0..3 {
        // A shorter segment fills only its first slots; the rest keep their old contents.
        let n = k.min(offsets[s + 1] - offsets[s]);
        println!("{:?} {:?}", &values[s * k..s * k + n], &positions[s * k..s * k + n]);
    }
}

Output

[0.9, 0.5] [0, 2]
[0.7] [0]
[0.8, 0.3] [1, 0]

Medians and quantiles over a sliding window

rolling_median(a, window) gives, at every position, the median of the last window values. It is pandas' Series.rolling(window).median(), with the same results. rolling_quantile(a, window, q) does the same for any quantile q between 0 and 1. A rolling median removes short spikes from a signal and keeps its steps sharp.

import numpy as np
import kwker

temps = np.array([20.5, 21.0, 35.0, 21.5, 22.0, 21.75, 22.5])  # (35.0: a sensor glitch)
print(kwker.rolling_median(temps, 3))
print(kwker.rolling_quantile(temps, 3, 0.75))

[  nan   nan 21.   21.5  22.   21.75 22.  ]
[   nan    nan 28.    28.25  28.5   21.875 22.25 ]

RustAdd kwker to Cargo.toml, then cargo run.

use kwker::Interpolation;

fn main() {
    let temps = [20.5, 21.0, 35.0, 21.5, 22.0, 21.75, 22.5]; // (35.0: a sensor glitch)
    let mut median = vec![0.0; temps.len()];
    kwker::rolling_median(&temps, 3, &mut median);
    println!("{median:?}");
    let mut q75 = vec![0.0; temps.len()];
    kwker::rolling_quantile(&temps, 3, 0.75, Interpolation::Linear, &mut q75);
    println!("{q75:?}");
}

Output

[NaN, NaN, 21.0, 21.5, 22.0, 21.75, 22.0]
[NaN, NaN, 28.0, 28.25, 28.5, 21.875, 22.25]

The interpolation argument picks how a quantile between two values is formed: "linear" (the default), "lower", "higher", "midpoint" or "nearest", as in pandas.

rolling_mad(a, window) gives the median absolute deviation of each window: the median of |x - m|, where m is the window's median. It measures how far values usually stray from the median, and one glitch barely moves it: a Hampel filter flags a value as a spike when it lies more than 3 x 1.4826 x MAD from its window's median.

import numpy as np
import kwker

temps = np.array([20.5, 21.0, 35.0, 21.5, 22.0, 21.75, 22.5])
mad = kwker.rolling_mad(temps, 3)
print(mad)
spike = np.abs(temps - kwker.rolling_median(temps, 3)) > 3 * 1.4826 * mad
print(np.flatnonzero(spike))

[ nan  nan 0.5  0.5  0.5  0.25 0.25]
[2]

RustAdd kwker to Cargo.toml, then cargo run.

fn main() {
    let temps = [20.5, 21.0, 35.0, 21.5, 22.0, 21.75, 22.5];
    let mut mad = vec![0.0; temps.len()];
    kwker::rolling_mad(&temps, 3, &mut mad);
    println!("{mad:?}");
}

Output

[NaN, NaN, 0.5, 0.5, 0.5, 0.25, 0.25]

Notes

Equal values: the one that came first is returned first, so results are the same on every run and every machine.
top_k_segments: a segment shorter than k fills its remaining positions with -1.
Rolling windows: the first window - 1 results are NaN, because the window is not full yet. A window that holds a NaN also gives NaN. Long windows cost little more than short ones.
1.4826 makes the MAD match the standard deviation of normally distributed data.

Sorting
Order and ranking
Statistics and data helpers: top-p, top-n-sigma and min-p filters of logits
API reference: top_k, select

Top-k and selection

The k largest or smallest values

Only the positions

The median and other positions

Sort only the beginning

Top-k with a filter

Top-k per group

Medians and quantiles over a sliding window

Notes

Related