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