Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- template <class T> using mpq = priority_queue<T, vector<T>, greater<T>>;
- class Solution {
- using ll = long long;
- using pii = pair<int, int>;
- vector<vector<ll>> dist;
- public:
- vector<vector<int>> modifiedGraphEdges(int n, vector<vector<int>>& edges, int source, int dest, int target) {
- constexpr int INF = INT_MAX / 2;
- dist = vector<vector<ll>>(1 << 7, vector<ll>(1 << 7, LLONG_MAX / 2));
- for (int i = 0; i < (1 << 7); ++i) dist[i][i] = 0;
- for (const auto& e: edges)
- dist[e[0]][e[1]] = (e[2] == -1 ? INF : e[2]), dist[e[1]][e[0]] = (e[2] == -1 ? INF : e[2]);
- vector<ll> ans(n + 1, LLONG_MAX / 2);
- vector<int> from(n + 1, -1);
- ans[source] = 0;
- /*
- for (int u = 1; u < n; ++u) {
- for (const auto& e: edges) {
- if (e[0] != u && e[1] != u) continue;
- int s = u, t = e[0] == u ? e[1] : e[0];
- if (ans[s] != LLONG_MAX / 2 && ans[t] > ans[s] + e[2])
- ans[t] = ans[s] + e[2], from[t] = s;
- }
- }
- */
- set<int> vis;
- using pli = pair<ll, int>;
- mpq<pli> q;
- q.push({0, source});
- while(q.empty() == false) {
- auto [d, u] = q.top(); q.pop();
- if (vis.count(u)) continue;
- vis.insert(u);
- ans[u] = min(ans[u], d);
- for (int v = 0; v < n; ++v) {
- ll w = dist[u][v];
- if (ans[v] > ans[u] + w) {
- ans[v] = ans[u] + w;
- from[v] = u;
- q.push({ans[v], v});
- }
- }
- }
- /*
- for (int k = 1; k <= n; ++k)
- for (int u = 1; u <= n; ++u)
- for (int v = 1; v <= n; ++v)
- for (int i = 0; i < n; ++i)
- cout << "ans[" << i << "]: " << ans[i] << endl;
- */
- if (ans[dest] < target) return {};
- if (ans[dest] == target) {
- for (auto& e: edges) {
- if (e[2] == -1) e[2] = INF;
- }
- return edges;
- }
- for (const auto& e: edges)
- if (e[2] == -1)
- dist[e[0]][e[1]] = dist[e[1]][e[0]] = 0;
- ans = vector<ll>(n + 1, LLONG_MAX / 2);
- vis.clear();
- q.push({0, source});
- while(q.empty() == false) {
- auto [d, u] = q.top(); q.pop();
- if (vis.count(u)) continue;
- vis.insert(u);
- ans[u] = min(ans[u], d);
- for (int v = 0; v < n; ++v) {
- ll w = dist[u][v];
- if (ans[v] > ans[u] + w) {
- ans[v] = ans[u] + w;
- from[v] = u;
- q.push({ans[v], v});
- }
- }
- }
- int block_cnt = 0;
- int now = dest;
- set<pii> tgt;
- while(now != source) {
- // cout << "now: " << now << " from: " << from[now] << endl;
- if (dist[now][from[now]] == 0) ++block_cnt, tgt.insert({min(now, from[now]), max(now, from[now])});
- now = from[now];
- }
- ll got = ans[dest];
- ll req = target - got, avg = (block_cnt != 0 ? req / block_cnt : 0);
- // cout << "block_cnt: " << block_cnt << endl;
- for (auto& e: edges) {
- // cout << "e: " << e[0] << ',' << e[1] << ',' << e[2] << endl;
- if (e[2] != -1) continue;
- int a = e[0], b = e[1];
- // if (from[a] != b && from[b] != a) continue;
- if (tgt.count({min(a, b), max(a, b)}) == 0) {
- e[2] = INF;
- continue;
- }
- e[2] = (block_cnt == 1 ? req : avg);
- req -= avg;
- block_cnt--;
- }
- return edges;
- }
- };
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement