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

import sys

line = sys.stdin.read().splitlines()
line_iter=iter(line)
input = lambda :next(line_iter)

def read(t=int):
    return list(map(t,input().split()))

N,M,K=read()
xs=[read() for _ in range(N)]
xt=read()

INF=10**18

def getdis(a,b)->int:
    res=0
    for l,r in zip(a,b):
        res+=(l-r)**2
    return res

def dimsum(a,b)->list[int]:
    res=[]
    for l,r in zip(a,b):
        res.append(l+r)
    return res

def dimdiv(a:list[int],b)->list[int]:
    res=[]
    for e in a:
        res.append(e//b)
    return res

def main():
    updated=True
    center=[xs[i] for i in range(K)]
    target=None # 每个i的元素对应的center编号

    while updated:
        new_target=[]
        center_avail=[N//K+1]*(N%K)+[N//K]*(K-N%K)

        for xidx,e in enumerate(xs):
            mindis=INF
            cidx=None
            for idx,c in enumerate(center):
                if center_avail[idx]<=0:
                    continue # 已经满了
                dis=getdis(e,c)
                if dis<mindis:
                    mindis=dis
                    cidx=idx
            assert cidx is not None,"1?"
            new_target.append(cidx)
            center_avail[cidx]-=1

        target=new_target

        new_center=[[0]*M for _ in range(K)]
        # update center
        for idx,e in enumerate(xs):
            cidx=target[idx]
            new_center[cidx] = dimsum(new_center[cidx],e)
        
        for i in range(N%K):
            new_center[i]=dimdiv(new_center[i],N//K+1)
        for i in range(N%K,K):
            new_center[i]=dimdiv(new_center[i],N//K)
        
        updated= center!=new_center or target!=new_target
        center=new_center
    
    
    center.sort()
    mindis=INF
    cidx=None
    for idx,c in enumerate(center):
        dis=getdis(xt,c)
        if dis<mindis:
            mindis=dis
            cidx=idx
    assert cidx is not None, "2?"
    for row in center:
        print(" ".join([str(e) for e in row]))
    print(cidx+1)
    
main()