Submission #676312

#TimeUsernameProblemLanguageResultExecution timeMemory
676312tengiz05Construction of Highway (JOI18_construction)C++17
100 / 100
428 ms30412 KiB
#include <bits/stdc++.h>

using i64 = long long;

constexpr int N = 1E5 + 5;

int a[N + 1];
void add(int x, int y) {
    for (int i = x + 1; i <= N; i += i & -i) {
        a[i] += y;
    }
}
int sum(int x) {
    int res = 0;
    for (int i = x; i > 0; i -= i & -i) {
        res += a[i];
    }
    return res;
}

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    
    int n;
    std::cin >> n;
    
    std::vector<int> a(n);
    for (int i = 0; i < n; i++) {
        std::cin >> a[i];
    }
    
    auto f = a;
    std::sort(f.begin(), f.end());
    f.erase(std::unique(f.begin(), f.end()), f.end());
    int m = f.size();
    for (int i = 0; i < n; i++) {
        a[i] = std::lower_bound(f.begin(), f.end(), a[i]) - f.begin();
    }
    
    std::vector<int> p(n, -1);
    
    std::vector<std::vector<int>> adj(n);
    
    std::vector<std::pair<int, int>> E;
    
    for (int _ = 0; _ < n - 1; _++) {
        int A, B;
        std::cin >> A >> B;
        A--;
        B--;
        adj[A].push_back(B);
        p[B] = A;
        E.emplace_back(A, B);
    }
    
    std::vector<int> siz(n, 1), pos(n), head(n);
    
    std::function<void(int)> dfs = [&](int u) {
        for (int &v : adj[u]) {
            dfs(v);
            if (siz[v] > siz[adj[u][0]]) {
                std::swap(v, adj[u][0]);
            }
            siz[u] += siz[v];
        }
    };
    dfs(0);
    
    int timeStamp = 0;
    
    std::function<void(int, int)> dfs2 = [&](int u, int h) {
        pos[u] = timeStamp++;
        head[u] = h;
        for (int v : adj[u]) {
            if (v == adj[u][0]) {
                dfs2(v, h);
            } else {
                dfs2(v, v);
            }
        }
    };
    dfs2(0, 0);
    
    std::vector<std::set<std::pair<int, int>>> s(n);
    for (int i = 0; i < n; i++) {
        s[head[i]].insert({pos[i], a[i]});
    }
    for (auto [A, B] : E) {
        int u = A;
        std::vector<std::pair<int, int>> v;
        while (u != -1) {
            int h = head[u];
            std::vector<std::pair<int, int>> cur;
            
            int last = pos[h] - 1;
            while (true) {
                assert(!s[h].empty());
                auto [r, c] = *s[h].begin();
                if (r <= pos[u]) {
                    cur.push_back({c, r - last});
                    s[h].erase(s[h].begin());
                    if (r == pos[u]) {
                        break;
                    }
                } else {
                    cur.push_back({c, pos[u] - last});
                    break;
                }
                last = r;
            }
            s[h].insert({pos[u], a[B]});
            
            std::reverse(cur.begin(), cur.end());
            v.insert(v.end(), cur.begin(), cur.end());
            u = p[h];
        }
        std::reverse(v.begin(), v.end());
        
        i64 ans = 0;
        for (auto [x, y] : v) {
            ans += y * 1LL * (sum(m) - sum(x + 1));
            add(x, y);
        }
        for (auto [x, y] : v) {
            add(x, -y);
        }
        std::cout << ans << "\n";
    }
    
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...