Submission #1129392

#	Time	Username	Problem	Language	Result	Execution time	Memory
1129392		patgra	Road Closures (APIO21_roads)	C++20	0 / 100	162 ms	55356 KiB

roads

#include <bits/stdc++.h>

#define rep(a,b,c) for(auto a = (b); a != (c); a++)
#define repD(a,b,c) for(auto a = (b); a != (c); a--)
#define repIn(a, b) for(auto& a : (b))
#define repIn2(a, b, c) for(auto& [a, b] : c)

#define ll long long

constexpr bool dbg = 0;
#define DEBUG if constexpr(dbg)
#define DC DEBUG std::cerr
#define eol std::endl

using namespace std;

constexpr int maxn = 1e5 + 7;
int n;
vector<pair<int, int>> graf[maxn];
vector<pair<ll, int>> tree[maxn];
int treeBase[maxn];
ll dp1[maxn], dp2[maxn], sumOfAll;
int vis[maxn];
unordered_map<int, int> whichEdge[maxn];
int k;
ll ans, baseAns;
bool deleted[maxn];
ll answers[maxn];

pair<ll, int> operator+(pair<ll, int> a, pair<ll, int> b) {
    return {a.first + b.first, a.second + b.second};
}

void treeAdd(int i, int x, int v) {
    // DC << "Tree[" << v << "][" << i << "] += " << x << eol;
    i += treeBase[v];
    tree[v][i] = {x, 1};
    i /= 2;
    while (i) {
        tree[v][i] = tree[v][2 * i] + tree[v][2 * i + 1];
        i /= 2;
    }
}
ll treeQuery(int x, int v) {
    if (x <= 0) return 0;
    DC << "Query " << v << " top " << x << eol;
    if (tree[v][1].second <= x) DC << " " << tree[v][1].first << eol;
    if (tree[v][1].second <= x) return tree[v][1].first;
    int i = 1;
    ll ret = 0;
    while (x) {
        if (tree[v][i * 2].second <= x) ret += tree[v][2 * i].first, x -= tree[v][2 * i].second, i = 2 * i + 1;
        else i = 2 * i;
    }
    DC << "  " << ret << eol;
    return ret;
}

void input() {
    scanf("%d", &n);
    int a, b, c;
    rep(i, 1, n) scanf("%d%d%d", &a, &b, &c), graf[a].emplace_back(b, c), graf[b].emplace_back(a, c);
}

queue<int> Q1, Q2;
void bfs() {
    Q1.push(0);
    while (!Q1.empty()) {
        auto v = Q1.front(); Q1.pop();
        Q2.push(v);
        vis[v] = -1;
        repIn2(u, c, graf[v]) if (!vis[u]) Q1.push(u);
    }
    swap(Q1, Q2);
}

bool cmp1(pair<int, int> a, pair<int, int> b) {
    return a.second > b.second;
}
bool cmp2(pair<int, int> a, pair<int, int> b) {
    return graf[a.first].size() > graf[b.first].size();
}

void init() {
    rep(v, 0, n) {
        ranges::sort(graf[v], cmp1);
        int xd = 0; repIn2(u, _, graf[v]) whichEdge[v][u] = xd++;
        ranges::sort(graf[v], cmp2);
        xd = graf[v].size() - 1;
        treeBase[v] = 1 << (32 - (xd == 0 ? 32 : __builtin_clz(xd)));
        treeBase[v] *= 2;
        tree[v].resize(treeBase[v] * 2);
        DC << "TreeBase[" << v << "] = " << treeBase[v] << eol;
    }
    DEBUG {
        DC << "Edges sorted by degree (descending order) : " << eol;
        rep(v, 0, n) {
            DC << " " << v << " ->  ";
            repIn2(u, c, graf[v]) DC << "(" << u << ' ' << c << ") ";
            DC << eol;
        }
        DC << "Edges sorted by cost (descending order) : " << eol;
        rep(v, 0, n) {
            DC << " " << v << " ->  ";
            repIn2(u, index, whichEdge[v]) DC << "(" << index << ": " << u << ") ";
            DC << eol;
        }
    }
    bfs();
    rep(v, 0, n) repIn2(u, c, graf[v]) sumOfAll += c;
    sumOfAll /= 2;
}

void dfs(int v) {
    vis[v] = k;
    dp1[v] = 0;
    int p;
    repIn2(u, c, graf[v]) {
        if (vis[u] == k) { p = u; continue; }
        if (graf[u].size() < k) break;
        dfs(u);
        dp1[v] += dp1[u];
    }
    dp2[v] = dp1[v];

    vector<ll> diffs;
    repIn2(u, c, graf[v]) {
        if (u == p) continue;
        if (graf[u].size() < k) break;
        diffs.push_back(dp2[u] + c - dp1[u]);
    }

    ranges::sort(diffs, greater<>());
    DEBUG {
        DC << "Diffs[" << v << "] = ";
        repIn(i, diffs) DC << i << ' ';
        DC << eol;
    }
    ll prefSum = 0, ogVal = dp1[v];
    rep(i, 0, k + 1) {
        dp1[v] = max(dp1[v], ogVal + prefSum + treeQuery(k - i, v));
        if (i == diffs.size()) break;
        prefSum += diffs[i];
    }
    prefSum = 0;
    rep(i, 0, k) {
        dp2[v] = max(dp2[v], ogVal + prefSum + treeQuery(k - i - 1, v));
        if (i == diffs.size()) break;
        prefSum += diffs[i];
    }
    DC << "DP[" << v << "] = " << dp1[v] << ' ' << dp2[v] << eol;
}

void addToNeighboursTrees(int v) {
    deleted[v] = true;
    repIn2(u, c, graf[v]) treeAdd(whichEdge[u][v], c, u);
    repIn2(u, c, graf[v]) if (deleted[u]) baseAns += c;
}

void calcForK() {
    while (!Q1.empty()) {
        auto v = Q1.front();
        Q1.pop();
        if (graf[v].size() < k) addToNeighboursTrees(v);
        else Q2.push(v);
    }
    ans = sumOfAll - baseAns;
    while (!Q2.empty()) {
        auto v = Q2.front();
        Q2.pop();
        Q1.push(v);
        if (vis[v] == k) continue;
        dfs(v);
        ans -= dp1[v];
    }
    // printf("%lld ", ans);
    answers[k] = ans;
}
void answer() {
    rep(i, 0, n) k = i, calcForK();
}

/*int main() {
    input();
    init();
    answer();
    printf("\n");
    return 0;
}*/

std::vector<long long> minimum_closure_costs(int N, std::vector<int> U,
                                             std::vector<int> V,
                                             std::vector<int> W) {
    n = N;
    int a, b, c;
    rep(i, 0, n - 1) a = U[i], b = W[i], c = W[i], graf[a].emplace_back(b, c), graf[b].emplace_back(a, c);
    init();
    answer();
    std::vector<ll> Vec;
    rep(i, 0, n) Vec.push_back(answers[i]);
    return Vec;
}

Compilation message (stderr)

roads.cpp: In function 'void input()':
roads.cpp:60:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   60 |     scanf("%d", &n);
      |     ~~~~~^~~~~~~~~~
roads.cpp:62:23: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   62 |     rep(i, 1, n) scanf("%d%d%d", &a, &b, &c), graf[a].emplace_back(b, c), graf[b].emplace_back(a, c);
      |                  ~~~~~^~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 0 / 5
• Subtask 2: 0 / 7
• Subtask 3: 0 / 14
• Subtask 4: 0 / 10
• Subtask 5: 0 / 17
• Subtask 6: 0 / 25
• Subtask 7: 0 / 22