Counting sort

Given is an array of n integers from the range 0..log(n). Sort the array.

Simple counting sort:

input A[_];
int n = A.size;
int k = log(n);

output int B[n];
int Cnt[k], ps[k], pos[k+1];

pardo(i : k) Cnt[i] = count(i, A);
prefix_sums(Cnt, ps);

pos[0]=0;
pardo(i:k) pos[i+1]=ps[i];

pardo(i : k) if (Cnt[i]>0) pardo(j : Cnt[i]) B[pos[i] + j] = i;
where the procedure count counts the number of occurences of i in A:
int count(int val, int X[_]) {
  int n = X.size, tmp[n];

  pardo(i : n)
    if (X[i] == val) tmp[i] = 1;
    else tmp[i] = 0;
  return sum(tmp);
}
Note that the previous algorithm has time O(log n) but the work is O(n log n). We improve on that; moreover, we will sort an array of key,value pairs according to the key values.
type kv { int key, val; }

input kv A[_];
int n = A.size, k = log(n);

#include "utils/prefix_sums.wt"

int B[k, n / k];

pardo(i : n / k) {
  pardo(j : k) B[j, i] = 0;
  for (int j = i * k; j < (i + 1) * k; j++) B[A[j].key, i]++;
}

pardo(i : k) {
  int tmp[n / k], ps[n / k];
  pardo(j : n / k) tmp[j] = B[i, j];
  prefix_sums(tmp, ps);
  pardo(j : n / k) B[i, j] = ps[j];
}

int start[k];
start[0] = 0;
for (int i = 1; i < k; i++) start[i] = start[i - 1] + B[i - 1, n / k - 1];

output kv S[n];

pardo(i : n / k) {
  int pos[k];
  pardo(j : k) {
    pos[j] = start[j];
    if (i > 0) pos[j] += B[j, i - 1];
  }
  for (int j = 0; j < k; j++) S[pos[A[i * k + j].key]++] = A[i * k + j];
}