均衡版 KMeans 分群与新用户归类

#include "bits/stdc++.h"
using namespace std;

int n,m,k,preLimit,otherLimit,n_mod_k;

struct guest{
    vector<int> features;
};

struct group{
    int centerID, id, limit = 0;
    vector<int> guestID;
    vector<int> center, preCenter;
    // 按center排序
    bool operator<(const group& o) const{
        return center < o.center;
    }
};

vector<guest> guests;
vector<group> groups;
vector<vector<int>> centers; // 存储中心

long long calcDis(vector<int> vector1, vector<int> &vector2) {
    long long res = 0;
    for(int i = 0; i < vector1.size(); i++){
        res += (vector1[i] - vector2[i]) * (vector1[i] - vector2[i]);
    }
    return res;
}

int main(){
    cin >> n >> m >> k;
    guests.resize(n);
    groups.resize(k);
    n_mod_k = n % k;
    if (n_mod_k == 0){
        preLimit = n / k;
        otherLimit = n / k;
    }else{
        preLimit = n / k + 1;
        otherLimit = n / k;
    }
    // 初试化
    for(int i = 0; i < n; i++){
        guests[i].features.resize(m);
        for(int j = 0; j < m; j++){
            cin >> guests[i].features[j]; // 保存特征
        }
        if (i < k) {
            groups[i].centerID = i;
            groups[i].guestID.push_back(i);
            groups[i].id = i;
            groups[i].preCenter = guests[i].features;
            groups[i].center = guests[i].features;
            //guests[i].centerID = i;
            if (i+1 <= n_mod_k) groups[i].limit = preLimit;
            else groups[i].limit = otherLimit;
            centers.push_back(guests[i].features);
        }else { // 为当前顾客找个中心组
            long long dis = 1e18;
            int tempGID = -1, preCenterID = -1;
            for(int j = 0; j < k; j++){
                group curG = groups[j];
                // 人数已满
                if (curG.guestID.size() >= curG.limit) continue;
                vector<int> v = guests[curG.centerID].features;
                long long temp = calcDis(guests[i].features, v);
                if (temp < dis){
                    dis = temp;
                    tempGID = curG.id;
                    preCenterID = groups[j].centerID;
                }else if(temp == dis){
                    if (preCenterID > groups[j].centerID){
                        tempGID = curG.id;
                        preCenterID = groups[j].centerID;
                    }
                }
            }
            // 将顾客加入当前组
            groups[tempGID].guestID.push_back(i);
        }
    }
    // 更新每组的中心
    while(true) {
        // 清空组中的顾客
        for(int i = 0; i < k; i++){
            groups[i].guestID.clear();
        }
        int count = 0;
        // 按顺序给顾客重新分配组
        for(int ii = 0; ii < n; ii++){
            vector<int> curGuest = guests[ii].features;
            long long dis = 1e18;
            int tempGID = -1, preCenterID = -1;
            for(int j = 0; j < k; j++){
                // 人数已满
                if (groups[j].guestID.size() >= groups[j].limit) continue;
                long long temp = calcDis(groups[j].center, curGuest);
                if (temp < dis){
                    dis = temp;
                    tempGID = groups[j].id;
                    preCenterID = groups[j].centerID;
                }else if(temp == dis){
                    if (preCenterID > groups[j].centerID){
                        tempGID = groups[j].id;
                        preCenterID = groups[j].centerID;
                    }
                }
            }
            // 将顾客加入当前组
            groups[tempGID].guestID.push_back(ii);
        }

        // 更新中心维度
        for (int i = 0; i < k; i++) {
            //group curG = groups[i];
            vector<int> vID = groups[i].guestID;
            // 更新m维特征
            for(int j = 0; j < m; j++){
                double tmp = 0.0;
                int sz = vID.size();
                for (int mm = 0; mm < sz; ++mm) {
                    vector<int> fe = guests[vID[mm]].features;
                    tmp += fe[j];
                }
                groups[i].center[j] = floor(tmp / sz);
            }
            if (!equal(groups[i].center.begin(),groups[i].center.end(),groups[i].preCenter.begin(),groups[i].preCenter.end())){
                groups[i].preCenter = groups[i].center;
            }else{
                count++;
            }
        }
        // 更新停止
        if (count == k) break;
    }

    // 将最终中心按字典序排序
    sort(groups.begin(),groups.end());
    // 读取最后一个新顾客的特征
    guest newGuest;
    newGuest.features.resize(m);
    for(int i = 0; i < m; i++){
        cin >> newGuest.features[i];
    }
    long long dis = 1e18;
    int tempID = -1;
    for(int j = 0; j < k; j++){
        group curG = groups[j];
        vector<int> v = newGuest.features;
        long long temp = calcDis(curG.center, v);
        if (temp < dis){
            dis = temp;
            tempID = j;
        }
    }
    // 输出中心
    for(int i = 0; i < k; i++){
        group curG = groups[i];
        vector<int> cen = curG.center;
        for(int j = 0; j < m; j++){
            cout << cen[j];
            if (j < m-1) cout << " ";
        }
        cout << "\n";
    }
    // 输出序号
    tempID++;
    cout << tempID << "\n";
    return 0;
}