UESTC 1146 秋实大哥与连锁快餐店 最小生成树、Prim

#include 
#include 
#include 
#include 
using namespace std;
const int maxn = 6669, INF = 0x3f3f3f3f;
int n;
//
struct point {
    int x, y, condi;
} po[maxn];
//
double dis[maxn];
bool v[maxn];
                                    //this problem is unique,so we don't use the vector g[maxn]

inline double getdist(point a,point b)
{
    return sqrt(pow(a.x-b.x,2) + pow(a.y-b.y,2));
}

double prim()
{
    memset(v, 0, sizeof v);
    for(int i = 1; i <= n; i++) dis[i] = INF;
    dis[1] = 0;
    double ans = 0;
    for(int i = 1; i <= n; i++){
        int mark = -1;
        for(int j = 1; j <= n; j++) if(!v[j]){
            if(mark == -1) mark = j;
            else if(dis[j] < dis) mark = j;
        }

        if(mark == -1) break;                       //all vertices are marked
        v = 1;
        ans += dis;

        for(int j = 1; j <= n; j++){
            if(j == mark) continue;
            if(!v[j]){
                dis[j] = min(dis[j], (po.condi==1 && po[j].condi==1) ? 0 : getdist(po,po[j]) );
            }
        }
    }
    return ans;
}

int main()
{
    scanf("%d", &n);
    for(int i = 1; i <= n; i++){
        scanf("%d%d%d", &po[i].x, &po[i].y, &po[i].condi);
    }
    printf("%.2f", prim());
    return 0;
}