1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128 | #include <iostream>
#include <vector>
#include <math.h>
#include <algorithm>
using namespace std;
#define lli long long int
#define vi vector<int>
#define vlli vector<lli>
#define vb vector<bool>
#define vvi vector<vi>
#define vvlli vector<vlli>
#define vpii vector<pair<int, int>>
#define vvpii vector<vpii>
#define str string
#define vs vector<str>
#define vc vector<char>
#define vvc vector<vc>
const lli mod = 1000000007;
void dfs(vvpii &T, vb &been, vlli &pathlens, int x, lli len)
{
if (been[x])
return;
been[x] = true;
if (len != 0)
pathlens.push_back(len);
for (int i = 0; i < T[x].size(); i++)
{
dfs(T, been, pathlens, T[x][i].first, len + T[x][i].second);
}
}
lli median(vlli arr, int l, int n)
{
sort(arr.begin() + l, arr.begin() + l + n);
return arr[n / 2];
}
int partition(vlli &arr, int l, int r, lli x)
{
int i;
for (i = l; i < r; i++)
if (arr[i] == x)
break;
swap(arr[i], arr[r]);
i = l;
for (int j = l; j <= r - 1; j++)
{
if (arr[j] <= x)
{
swap(arr[i], arr[j]);
i++;
}
}
swap(arr[i], arr[r]);
return i;
}
int kthSmallest(vlli &arr, int l, int r, lli k)
{
int n = r - l + 1;
int i;
vlli med((n + 4) / 5);
for (i = 0; i < n / 5; i++)
med[i] = median(arr, l + i * 5, 5);
if (i * 5 < n)
{
med[i] = median(arr, l + i * 5, n % 5);
i++;
}
lli medOfMed = (i == 1) ? med[i - 1] : kthSmallest(med, 0, i - 1, i / 2);
int pos = partition(arr, l, r, medOfMed);
if (pos - l == k - 1)
return arr[pos];
if (pos - l > k - 1)
return kthSmallest(arr, l, pos - 1, k);
return kthSmallest(arr, pos + 1, r, k - pos + l - 1);
}
int main()
{
int n, k;
cin >> n >> k;
vvpii T(n + 1);
vlli powers;
powers.push_back(1);
for (int i = 0; i < n - 1; i++)
{
int a, b, p;
cin >> a >> b >> p; // p <= 4
while (powers.size() < p + 1)
{
lli x = powers[powers.size() - 1];
x *= n;
x = x % mod;
powers.push_back(x);
}
lli w = powers[p];
T[a].push_back(make_pair(b, w));
T[b].push_back(make_pair(a, w));
}
vlli pathlens;
for (int i = 1; i <= n; i++)
{
vb been(n + 1, false);
lli len = 0;
dfs(T, been, pathlens, i, len);
}
cout << kthSmallest(pathlens, 0, pathlens.size() - 1, 2 * k) << endl;
return 0;
}
// sort(pathlens.begin(), pathlens.end());
// for (int i = 0; i < pathlens.size(); i++) // DEBUG
// cout << pathlens[i] << " ";
// cout << endl;
|