2698232_ac_200ms_1528k.cpp

来自「北大大牛代码 1240道题的原代码 超级权威」· C++ 代码 · 共 106 行

#include <stdio.h>
#include <stdlib.h>
#include <algorithm>
#include <string.h>
#define INIT (Tree *)malloc(sizeof(Tree))
using namespace std;

typedef struct node
{
	int id;
	struct node *next;
}Tree;

typedef struct Node
{
	Tree *head;
}T;

T tree[10001];

struct NODE
{
	int st, ed;
	__int64 tot;
	__int64 w, id;
}Edge[10000];

bool cmp1(struct NODE a,struct NODE b)
{
	return a.tot > b.tot;
}

bool cmp2(struct NODE a,struct NODE b)
{
	return a.tot < b.tot;
}

int n, k, sh, sc;
int visited[10001];
__int64 num[10001];
int grade[10001];

__int64 dfs(int v,int g)
{
	Tree *p;
	
	visited[v] =1;
	grade[v] = g;
	p = tree[v].head;
	while(p->next)
	{
		if(!visited[p->next->id])
			num[v] += dfs(p->next->id,g+1);
		p = p->next;
	}
	return num[v];
}

int main()
{
	int i;
	Tree *p;
	__int64 t, w;
	int st, ed;
	
	scanf("%d%d%d%d",&n,&k,&sh,&sc);
	for(i = 0; i < n; i++)
	{
		num[i] = 1;
		tree[i].head = INIT;
		tree[i].head->next = NULL;
	}
	for(i = 0; i < n-1; i++)
	{
		scanf("%d%d%I64d",&st,&ed,&w);
		Edge[i].id = i+1;
		Edge[i].w = w;
		Edge[i].st = st-1;Edge[i].ed = ed-1;
		p = INIT;
		p->id = ed-1;
		p->next = tree[st-1].head->next;
		tree[st-1].head->next = p;
		p = INIT;
		p->id = st-1;
		p->next = tree[ed-1].head->next;
		tree[ed-1].head->next = p;
	}
	memset(visited,0,sizeof(visited));
	dfs(0,0);
	t = n;int nn;
	for(i = 0; i < n-1; i++)
	{
		if(grade[Edge[i].st]>grade[Edge[i].ed])
			nn = Edge[i].st;
		else
			nn = Edge[i].ed;
		Edge[i].tot = Edge[i].w*num[nn]*(t-num[nn]);
	}
	if(sc>=sh)
		sort(Edge,Edge+n-1,cmp1);
	else
		sort(Edge,Edge+n-1,cmp2);
	for(i = 0; i < k; i++)
		printf("%d ",Edge[i].id);
	return 0;
}