第32章 Top-K和KNN

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

#define DIM 2
#define MAX_STUDENTS 100
#define NAME_LEN 32

struct Student {
    int id;
    char name[NAME_LEN];
    float vec[DIM];
};

struct StudentDB {
    struct Student rows[MAX_STUDENTS];
    int count;
};

struct Result {
    int index;      // 在数据库中的下标
    float dist;     // 到查询点的距离
};

void db_init(struct StudentDB *db)
{
    // 将记录数置零，数据库初始化为空
    db->count = 0;
}

int db_insert(struct StudentDB *db, int id, const char *name, const float vec[DIM])
{
    // 数据库已满则插入失败
    if (db->count >= MAX_STUDENTS) return 0;
    // 取末尾位置写入新记录
    struct Student *s = &db->rows[db->count];
    s->id = id;
    snprintf(s->name, sizeof(s->name), "%s", name);
    // 逐维度复制向量分量
    for (int i = 0; i < DIM; i++) s->vec[i] = vec[i];
    // 记录数加 1
    db->count++;
    return 1;
}

float euclidean_dist(const float a[DIM], const float b[DIM])
{
    float sum = 0.0f;
    for (int i = 0; i < DIM; i++) {
        // 各维度差值的平方累加
        float d = a[i] - b[i];
        sum += d * d;
    }
    // 开平方得到欧氏距离
    return sqrtf(sum);
}

// 用于 qsort：按距离从小到大
int cmp_by_dist(const void *a, const void *b)
{
    float da = ((const struct Result *)a)->dist;
    float db = ((const struct Result *)b)->dist;
    if (da < db) return -1;
    if (da > db) return 1;
    return 0;
}

int knn_sort(const struct StudentDB *db, const float query[DIM],
             int k, struct Result *out)
{
    // 存放所有记录的(下标, 距离)
    struct Result all[MAX_STUDENTS];
    int n = db->count;

    // 参数无效时直接返回 0
    if (k <= 0 || n <= 0) {
        return 0;
    }

    // 遍历数据库，计算每条记录到查询点的距离
    for (int i = 0; i < n; i++) {
        all[i].index = i;
        all[i].dist = euclidean_dist(query, db->rows[i].vec);
    }

    // 按距离从小到大排序
    qsort(all, n, sizeof(struct Result), cmp_by_dist);

    // 取前 K 个（记录不足时取全部）
    int found = (k < n) ? k : n;
    for (int i = 0; i < found; i++) {
        out[i] = all[i];
    }
    return found;
}

int main(void)
{
    struct StudentDB db;
    // 初始化数据库
    db_init(&db);

    // 插入 6 条学生记录
    float v1[] = {21.0f, 0.95f};  db_insert(&db, 1, "Alice", v1);
    float v2[] = {5.0f,  0.60f};  db_insert(&db, 2, "Bob",   v2);
    float v3[] = {15.0f, 0.85f};  db_insert(&db, 3, "Carol", v3);
    float v4[] = {18.0f, 0.90f};  db_insert(&db, 4, "Dave",  v4);
    float v5[] = {8.0f,  0.65f};  db_insert(&db, 5, "Eve",   v5);
    float v6[] = {23.0f, 0.92f};  db_insert(&db, 6, "Frank", v6);

    // 查询向量和 K 值
    float query[DIM] = {16.0f, 0.88f};
    int k = 3;

    // 执行 KNN 查询
    struct Result results[MAX_STUDENTS];
    int found = knn_sort(&db, query, k, results);

    // 打印结果
    printf("Query point: [%.2f, %.2f]\n", query[0], query[1]);
    printf("Top-%d nearest neighbors:\n", k);
    for (int i = 0; i < found; i++) {
        const struct Student *s = &db.rows[results[i].index];
        printf("  %d. %-8s  dist=%.2f\n", i + 1, s->name, results[i].dist);
    }

    return 0;
}

// 找 results 里距离最大的那个的下标
int find_max_index(const struct Result *results, int count)
{
    int max_i = 0;
    for (int i = 1; i < count; i++) {
        if (results[i].dist > results[max_i].dist) {
            max_i = i;
        }
    }
    return max_i;
}

int knn_stream(const struct StudentDB *db, const float query[DIM],
               int k, struct Result *out)
{
    int n = db->count;

    // 参数无效时直接返回 0
    if (k <= 0 || n <= 0) {
        return 0;
    }

    int cap = (k < n) ? k : n;  // 实际能取的数量
    int filled = 0;

    // 逐条扫描数据库
    for (int i = 0; i < n; i++) {
        float d = euclidean_dist(query, db->rows[i].vec);

        if (filled < cap) {
            // 还没填满，直接加进去
            out[filled].index = i;
            out[filled].dist = d;
            filled++;
        } else {
            // 已满，看是否比最远的更近
            int worst = find_max_index(out, cap);
            if (d < out[worst].dist) {
                // 替换掉当前最远的那条
                out[worst].index = i;
                out[worst].dist = d;
            }
        }
    }
    return filled;
}

int found = knn_stream(&db, query, k, results);

void db_knn(const struct StudentDB *db, const float query[DIM], int k)
{
    // 用流式算法找出 Top-K
    struct Result results[MAX_STUDENTS];
    int found = knn_stream(db, query, k, results);

    // 结果按距离从小到大排序
    qsort(results, found, sizeof(results[0]), cmp_by_dist);

    // 打印排序后的结果
    printf("Query point: [%.2f, %.2f]  Top-%d:\n", query[0], query[1], k);
    for (int i = 0; i < found; i++) {
        const struct Student *s = &db->rows[results[i].index];
        printf("  %d. id=%d %-8s dist=%.2f\n",
               i + 1, s->id, s->name, results[i].dist);
    }
}

db_knn(&db, query, 5);