水塘抽样：未知大小数据流的公平抽样

#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <assert.h>

/* 生成 [lo, hi] 之间的均匀随机整数 */
static long rand_range(long lo, long hi)
{
    assert(hi >= lo);
    long range = hi - lo + 1;
    /* 拒绝采样消除模偏差 */
    long limit = RAND_MAX - (RAND_MAX % range);
    long r;
    do {
        r = random();
    } while (r >= limit);
    return lo + r % range;
}

/*
 * Algorithm R — 经典水塘抽样
 *
 * reservoir: 预分配的大小为 k 的数组
 * stream:    输入数据流数组（模拟）
 * n:         流的长度
 * k:         样本大小
 *
 * 返回值: 实际采样数（min(n, k)）
 */
int reservoir_sample_R(int *reservoir, const int *stream, int n, int k)
{
    if (n <= 0 || k <= 0)
        return 0;

    int actual_k = (n < k) ? n : k;

    /* 阶段一：填充 */
    for (int i = 0; i < actual_k; i++)
        reservoir[i] = stream[i];

    /* 阶段二：替换 */
    for (int i = k; i < n; i++) {
        long j = rand_range(0, i);  /* [0, i] 均匀随机 */
        if (j < k)
            reservoir[j] = stream[i];
    }

    return actual_k;
}

#include <math.h>

/* 生成 (0, 1) 开区间的均匀随机浮点数 */
static double rand_uniform(void)
{
    double r;
    do {
        r = (double)random() / (double)RAND_MAX;
    } while (r == 0.0 || r == 1.0);
    return r;
}

/*
 * Algorithm L — 跳跃优化水塘抽样
 *
 * 与 Algorithm R 接口一致，但内部通过跳跃减少随机数生成次数。
 */
int reservoir_sample_L(int *reservoir, const int *stream, int n, int k)
{
    if (n <= 0 || k <= 0)
        return 0;

    int actual_k = (n < k) ? n : k;

    /* 阶段一：填充 */
    for (int i = 0; i < actual_k; i++)
        reservoir[i] = stream[i];

    if (n <= k)
        return actual_k;

    /* 初始化权重 W */
    double W = exp(log(rand_uniform()) / k);

    int i = k;
    while (1) {
        /* 计算跳跃距离 */
        double skip_d = floor(log(rand_uniform()) / log(1.0 - W));
        /* 防止 skip 溢出 */
        if (skip_d > (double)(n - i))
            break;
        int skip = (int)skip_d;
        i += skip + 1;
        if (i >= n)
            break;
        /* 替换水塘中随机位置 */
        int j = (int)(rand_uniform() * k);
        if (j >= k) j = k - 1;  /* 安全边界 */
        reservoir[j] = stream[i];
        /* 更新权重 */
        W *= exp(log(rand_uniform()) / k);
    }

    return actual_k;
}

typedef struct {
    int    value;
    double key;
} weighted_item_t;

/* 最小堆操作：维护堆顶为最小键 */
static void heap_sift_down(weighted_item_t *heap, int size, int pos)
{
    while (1) {
        int smallest = pos;
        int left  = 2 * pos + 1;
        int right = 2 * pos + 2;
        if (left < size && heap[left].key < heap[smallest].key)
            smallest = left;
        if (right < size && heap[right].key < heap[smallest].key)
            smallest = right;
        if (smallest == pos) break;
        weighted_item_t tmp = heap[pos];
        heap[pos] = heap[smallest];
        heap[smallest] = tmp;
        pos = smallest;
    }
}

static void heap_build(weighted_item_t *heap, int size)
{
    for (int i = size / 2 - 1; i >= 0; i--)
        heap_sift_down(heap, size, i);
}

/*
 * A-Res 加权水塘抽样
 *
 * values:  元素值数组
 * weights: 对应权重数组
 * n:       流长度
 * k:       样本大小
 * result:  输出的样本值数组（大小 k）
 */
int weighted_reservoir_ares(int *result, const int *values,
                            const double *weights, int n, int k)
{
    if (n <= 0 || k <= 0)
        return 0;

    int actual_k = (n < k) ? n : k;
    weighted_item_t *heap = malloc(k * sizeof(weighted_item_t));
    if (!heap) return 0;

    /* 阶段一：填充 */
    for (int i = 0; i < actual_k; i++) {
        heap[i].value = values[i];
        heap[i].key   = pow(rand_uniform(), 1.0 / weights[i]);
    }

    if (n <= k) {
        for (int i = 0; i < actual_k; i++)
            result[i] = heap[i].value;
        free(heap);
        return actual_k;
    }

    /* 建最小堆 */
    heap_build(heap, k);

    /* 阶段二：替换 */
    for (int i = k; i < n; i++) {
        double key = pow(rand_uniform(), 1.0 / weights[i]);
        if (key > heap[0].key) {
            heap[0].value = values[i];
            heap[0].key   = key;
            heap_sift_down(heap, k, 0);
        }
    }

    for (int i = 0; i < k; i++)
        result[i] = heap[i].value;
    free(heap);
    return k;
}

/*
 * A-ExpJ 加权水塘抽样（跳跃优化）
 * 简化实现，展示核心思路
 */
int weighted_reservoir_aexpj(int *result, const int *values,
                             const double *weights, int n, int k)
{
    if (n <= 0 || k <= 0)
        return 0;

    int actual_k = (n < k) ? n : k;
    weighted_item_t *heap = malloc(k * sizeof(weighted_item_t));
    if (!heap) return 0;

    /* 阶段一：填充 */
    for (int i = 0; i < actual_k; i++) {
        heap[i].value = values[i];
        heap[i].key   = pow(rand_uniform(), 1.0 / weights[i]);
    }

    if (n <= k) {
        for (int i = 0; i < actual_k; i++)
            result[i] = heap[i].value;
        free(heap);
        return actual_k;
    }

    heap_build(heap, k);

    /* 阶段二：指数跳跃替换 */
    double threshold = heap[0].key;
    double weight_sum = 0.0;
    double jump_target = log(rand_uniform()) / log(threshold);

    int i = k;
    while (i < n) {
        weight_sum += weights[i];
        if (weight_sum >= jump_target) {
            /* 此元素被选中，计算条件键 */
            double tw = pow(threshold, weights[i]);
            double key = pow(tw + rand_uniform() * (1.0 - tw),
                            1.0 / weights[i]);
            /* 替换堆顶 */
            heap[0].value = values[i];
            heap[0].key   = key;
            heap_sift_down(heap, k, 0);
            /* 更新阈值和跳跃目标 */
            threshold = heap[0].key;
            weight_sum = 0.0;
            jump_target = log(rand_uniform()) / log(threshold);
        }
        i++;
    }

    for (int i = 0; i < k; i++)
        result[i] = heap[i].value;
    free(heap);
    return k;
}

typedef struct {
    int *reservoir;     /* 水塘数组 */
    int  k;             /* 水塘大小 */
    long count;         /* 已处理元素数 */
} reservoir_state_t;

/*
 * 合并两个水塘
 * 结果存入 dst，src 的内容被消费
 */
void reservoir_merge(reservoir_state_t *dst, const reservoir_state_t *src)
{
    long total = dst->count + src->count;

    /* 对于 src 中的每个元素，以 src->count / total 的概率纳入 */
    for (int i = 0; i < src->k; i++) {
        /* 元素在 src 的水塘中的"等效概率"是 1/src->count × src->k */
        /* 简化方案：从合并的 2k 个候选中随机选 k 个 */
        long j = rand_range(0, total - 1);
        if (j < src->count) {
            /* 用 src 的元素替换 dst 中的随机位置 */
            int pos = (int)rand_range(0, dst->k - 1);
            dst->reservoir[pos] = src->reservoir[i];
        }
    }

    dst->count = total;
}

/*
 * 更精确的合并：两个水塘按比例采样
 */
void reservoir_merge_exact(int *result, int k,
                           const int *r1, long n1,
                           const int *r2, long n2)
{
    long total = n1 + n2;
    int idx = 0;

    /* 从 r1 中以 n1/total 的期望比例采样 */
    for (int i = 0; i < k && idx < k; i++) {
        if (rand_range(1, total - idx) <= (long)(k * n1 / total) - idx / 2) {
            result[idx++] = r1[i];
        }
    }
    /* 从 r2 中补齐 */
    for (int i = 0; idx < k; i++) {
        result[idx++] = r2[i % k];
    }
}

/*
 * reservoir_sampling.c — 水塘抽样完整实现
 *
 * 包含：Algorithm R, Algorithm L, A-Res, 分布式合并, 验证测试
 * 编译：gcc -O2 -o reservoir reservoir_sampling.c -lm
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <assert.h>

/* ========== 随机数工具 ========== */

static long rand_range(long lo, long hi)
{
    assert(hi >= lo);
    long range = hi - lo + 1;
    long limit = RAND_MAX - (RAND_MAX % range);
    long r;
    do {
        r = random();
    } while (r >= limit);
    return lo + r % range;
}

static double rand_uniform(void)
{
    double r;
    do {
        r = (double)random() / (double)RAND_MAX;
    } while (r == 0.0 || r == 1.0);
    return r;
}

/* ========== Algorithm R ========== */

int reservoir_R(int *reservoir, const int *stream, int n, int k)
{
    if (n <= 0 || k <= 0) return 0;
    int actual = (n < k) ? n : k;

    for (int i = 0; i < actual; i++)
        reservoir[i] = stream[i];

    for (int i = k; i < n; i++) {
        long j = rand_range(0, i);
        if (j < k)
            reservoir[j] = stream[i];
    }
    return actual;
}

/* ========== Algorithm L ========== */

int reservoir_L(int *reservoir, const int *stream, int n, int k)
{
    if (n <= 0 || k <= 0) return 0;
    int actual = (n < k) ? n : k;

    for (int i = 0; i < actual; i++)
        reservoir[i] = stream[i];

    if (n <= k) return actual;

    double W = exp(log(rand_uniform()) / k);
    int i = k - 1;

    while (1) {
        double skip = floor(log(rand_uniform()) / log(1.0 - W));
        if (skip > (double)(n - i - 1)) break;
        i += (int)skip + 1;
        if (i >= n) break;
        reservoir[(int)(rand_uniform() * k) % k] = stream[i];
        W *= exp(log(rand_uniform()) / k);
    }
    return actual;
}

/* ========== k=1 特例 ========== */

int reservoir_single(const int *stream, int n)
{
    assert(n > 0);
    int result = stream[0];
    for (int i = 1; i < n; i++) {
        if (rand_range(0, i) == 0)
            result = stream[i];
    }
    return result;
}

/* ========== A-Res 加权采样 ========== */

typedef struct {
    int    value;
    double key;
} witem_t;

static void min_heap_sift(witem_t *h, int size, int pos)
{
    while (1) {
        int s = pos, l = 2*pos+1, r = 2*pos+2;
        if (l < size && h[l].key < h[s].key) s = l;
        if (r < size && h[r].key < h[s].key) s = r;
        if (s == pos) break;
        witem_t t = h[pos]; h[pos] = h[s]; h[s] = t;
        pos = s;
    }
}

static void min_heap_build(witem_t *h, int size)
{
    for (int i = size/2 - 1; i >= 0; i--)
        min_heap_sift(h, size, i);
}

int reservoir_ares(int *result, const int *vals, const double *wts,
                   int n, int k)
{
    if (n <= 0 || k <= 0) return 0;
    int actual = (n < k) ? n : k;
    witem_t *heap = malloc(k * sizeof(witem_t));
    if (!heap) return 0;

    for (int i = 0; i < actual; i++) {
        heap[i].value = vals[i];
        heap[i].key   = pow(rand_uniform(), 1.0 / wts[i]);
    }

    if (n <= k) {
        for (int i = 0; i < actual; i++) result[i] = heap[i].value;
        free(heap);
        return actual;
    }

    min_heap_build(heap, k);

    for (int i = k; i < n; i++) {
        double key = pow(rand_uniform(), 1.0 / wts[i]);
        if (key > heap[0].key) {
            heap[0].value = vals[i];
            heap[0].key   = key;
            min_heap_sift(heap, k, 0);
        }
    }

    for (int i = 0; i < k; i++) result[i] = heap[i].value;
    free(heap);
    return k;
}

/* ========== 验证函数 ========== */

/*
 * 验证均匀性：对每个元素统计被选中的次数，
 * 理论期望 = trials * k / n
 */
void verify_uniformity(const char *name,
                       int (*sampler)(int*, const int*, int, int),
                       int n, int k, int trials)
{
    int *stream = malloc(n * sizeof(int));
    int *reservoir = malloc(k * sizeof(int));
    long *count = calloc(n, sizeof(long));

    for (int i = 0; i < n; i++)
        stream[i] = i;

    for (int t = 0; t < trials; t++) {
        sampler(reservoir, stream, n, k);
        for (int j = 0; j < k; j++)
            count[reservoir[j]]++;
    }

    double expected = (double)trials * k / n;
    double max_dev = 0;
    for (int i = 0; i < n; i++) {
        double dev = fabs((double)count[i] - expected) / expected;
        if (dev > max_dev) max_dev = dev;
    }

    printf("[%s] n=%d, k=%d, trials=%d\n", name, n, k, trials);
    printf("  期望次数: %.1f, 最大偏差: %.2f%%\n",
           expected, max_dev * 100);
    printf("  结论: %s\n\n",
           max_dev < 0.05 ? "通过（偏差 < 5%%）" : "可能存在偏差");

    free(stream);
    free(reservoir);
    free(count);
}

/* ========== 性能基准 ========== */

void benchmark(const char *name,
               int (*sampler)(int*, const int*, int, int),
               int n, int k)
{
    int *stream = malloc(n * sizeof(int));
    int *reservoir = malloc(k * sizeof(int));

    for (int i = 0; i < n; i++)
        stream[i] = i;

    struct timespec start, end;
    clock_gettime(CLOCK_MONOTONIC, &start);

    int rounds = 100;
    for (int t = 0; t < rounds; t++)
        sampler(reservoir, stream, n, k);

    clock_gettime(CLOCK_MONOTONIC, &end);

    double elapsed = (end.tv_sec - start.tv_sec)
                   + (end.tv_nsec - start.tv_nsec) / 1e9;
    printf("[%s] n=%d, k=%d, %d 轮: %.3f ms/轮\n",
           name, n, k, rounds, elapsed / rounds * 1000);

    free(stream);
    free(reservoir);
}

/* ========== 主程序 ========== */

int main(void)
{
    srandom(time(NULL));

    printf("===== 均匀性验证 =====\n\n");
    verify_uniformity("Algorithm R", reservoir_R, 100, 10, 100000);
    verify_uniformity("Algorithm L", reservoir_L, 100, 10, 100000);

    printf("===== 性能基准 =====\n\n");
    benchmark("Algorithm R", reservoir_R, 1000000, 100);
    benchmark("Algorithm L", reservoir_L, 1000000, 100);

    printf("===== k=1 特例测试 =====\n\n");
    {
        int stream[] = {0, 1, 2, 3, 4};
        int counts[5] = {0};
        int trials = 100000;
        for (int t = 0; t < trials; t++)
            counts[reservoir_single(stream, 5)]++;
        printf("k=1, n=5, %d 轮:\n", trials);
        for (int i = 0; i < 5; i++)
            printf("  元素 %d: %.2f%% (期望 20.00%%)\n",
                   i, 100.0 * counts[i] / trials);
    }

    return 0;
}

#include <pthread.h>

typedef struct {
    int    *reservoir;
    int     k;
    long    count;
    /* 每个实例独立的 RNG 状态 */
    unsigned long rng_state[4];  /* xoshiro256** */
} thread_safe_reservoir_t;

/* 从流中均匀随机选一个元素 */
int sample_one(stream_t *s)
{
    int result = stream_next(s);
    int count = 1;
    while (stream_has_next(s)) {
        count++;
        int x = stream_next(s);
        if (rand_range(1, count) == 1)
            result = x;
    }
    return result;
}