#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;
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; |
English