一直好奇向量数据库的索引是如何实现的，我们可以推断向量搜索的简单实现：把数据存入向量数据库时，会计算每个分段文档的向量(文档向量），然后把分段文档和文档向量同时存入向量数据库。从向量数据库中搜索文档时，会把待搜索问题转为向量(问题向量），然后计算问题向量与所有文档向量的距离，数据库会返回距离最短的一个或多个文档。

上面计算方式，返回的文档与问题最相关，但文档多时，耗费的计算资源同样多，有没有更好的方法？最容易想到的方法是把存入数据库中的文档向量先聚类为 K 个簇，从向量数据库中搜索文档时，先找到最相近的簇，再和簇内的每个文档向量比较找到与问题向量距离最短的一个或多个文档。这样需要的计算量就会少很多。

K-means算法简介

怎么把文档向量聚类为 K 个簇？可以使用K-means 算法。

1. 初始化：随机选择K个数据点作为初始簇中心 。
2. 分配样本：计算每个数据点到各簇中心的距离（通常用欧氏距离），将其分配到最近的簇 。
3. 更新簇中心：重新计算每个簇的均值作为新中心 。
4. 迭代：重复分配和更新步骤，直到中心不再变化或达到最大迭代次数

K-means 算法 python 实现

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs

# 生成模拟数据，300个数据点，4个簇，每个簇的标准差为0.6，随机种子为0
X, _ = make_blobs(n_samples=300, centers=4, cluster_std=0.60, random_state=0)

# K-means算法实现
def k_means(X, K, max_iters=100):
    np.random.seed(42)
    # 1. 随机初始化：从数据集中随机选择K个点作为初始簇中心
    # np.random.choice 是随机抽样函数
    # X.shape[0] 表示模拟数据个数；
    # replace=False 表示无放回抽样，即每个数据只能被选中一次，不允许重复
    centroids = X[np.random.choice(X.shape[0], K, replace=False)]
    original_centroids = centroids.copy()

    for iteration in range(max_iters):
        # 2. 分配数据点：计算每个数据点到各个簇中心的距离
        # X[:, np.newaxis] 表示为 X 添加一个新维度，使其形状变为 (n_samples, 1, n_features)。
        # 目的是通过广播机制与 centroids 的维度 (K, n_features) 对齐，便于后续逐元素计算
        # X[:, np.newaxis] - centroids 广播机制会将 X 扩展为 (n_samples, K, n_features)，
        # centroids 扩展为 (n_samples, K, n_features)（实际仅逻辑扩展，不占用额外内存）
        distances = np.linalg.norm(X[:, np.newaxis] - centroids, axis=2)
        # 3. 将每个数据点分配到最近的簇
        labels = np.argmin(distances, axis=1)
 
        # 4. 更新簇中心：重新计算每个簇的中心点，即簇内所有点的均值
        # labels == k 生成布尔掩码（Boolean Mask），筛选出 labels 数组中标签等于 k 的所有样本索引
        # X[labels == k] 根据布尔掩码从数据矩阵 X 中提取属于第 k 个簇的所有样本
        # .mean(axis=0) 沿列方向（axis=0）计算均值，得到第 k 个簇的中心坐标
        new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(K)])
      
        # 5. 迭代优化：检查簇中心是否变化
        if np.all(centroids == new_centroids):
            print(f"算法在第{iteration + 1}次迭代后收敛")
            break
        centroids = new_centroids

    return original_centroids, centroids, labels


# 使用K-means算法进行聚类
K = 4  # 指定簇的数量
original_centroids, centroids, labels = k_means(X, K)
# 结果可视化
plt.figure(figsize=(8, 6))
# 每个簇的数据使用不同颜色展示
plt.scatter(X[:, 0], X[:, 1], c=labels, s=50, cmap='viridis', edgecolor='k')
# 最终得到的簇心用红色展示
plt.scatter(centroids[:, 0], centroids[:, 1], c='red', s=200, alpha=0.5, marker='X')
# 最初得到的簇心用蓝色展示
plt.scatter(origcentroids[:, 0], origcentroids[:, 1], c='blue', s=200, alpha=0.5, marker='X')
plt.title("K-means Clustering")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
# 保存到本地图片
plt.savefig("k_means_clustering.png")

上面实现中 distances = np.linalg.norm(X[:, np.newaxis] - centroids, axis=2) 等价于下面代码，计算每个数据点到各个簇中心的距离。

distances = np.zeros((X.shape[0], centroids.shape[0]))
for i in range(X.shape[0]):
    for j in range(centroids.shape[0]):
        distances[i, j] = np.sqrt(np.sum((X[i] - centroids[j])**2))

展示最终得到的簇中心与各个簇如下图所示

同时展示 K-means 算法最初随机选择的簇中心，使用蓝色叉展示如下，与最终的簇中心位置相隔较远，可见算法帮我们动态调整了簇中心。

向量数据库进行实际的向量搜索时，并不是简单通过 K-meams 算法创建索引，本文是从一个算法小白角度推导向量搜索如何实现，接下来会逐步介绍实际使用的向量搜索算法。