Submission #485355

#	Time	Username	Problem	Language	Result	Execution time	Memory
485355		Lam_lai_cuoc_doi	Transport (COCI19_transport)	C++17	130 / 130	417 ms	25084 KiB

Main

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using ld = long double;
using ull = unsigned long long;

constexpr bool typetest = 0;
constexpr int N = 2e5 + 5;

struct FenwickTree
{
    int n;
    int a[N];
    FenwickTree()
    {
        memset(a, 0, sizeof a);
    }
    void Assign(int n)
    {
        this->n = n;
        fill(a, a + n + 1, 0);
    }
    void Update(int p, int v)
    {
        for (; p <= n; p += p & -p)
            a[p] += v;
    }
    int Get(int p)
    {
        int ans(0);
        for (; p; p -= p & -p)
            ans += a[p];
        return ans;
    }
    int Get(int l, int r)
    {
        return Get(r) - Get(l - 1);
    }
} fu, fd;

int n, cnt[N], par[N];
bool inused[N];
ll du[N], dd[N], maxn[N];
ll w[N], a[N], ans(0);
vector<ll> su, sd;
vector<pair<int, int>> adj[N];

void Read()
{
    cin >> n;

    for (int i = 1; i <= n; ++i)
        cin >> a[i];

    for (int i = 1; i < n; ++i)
    {
        inused[i] = 1;
        int u, v;
        cin >> u >> v >> w[i];
        adj[v].emplace_back(u, i);
        adj[u].emplace_back(v, i);
    }
}

void dfs_cnt(int v, int p = -1)
{
    cnt[v] = 1;
    for (auto i : adj[v])
        if (i.first != p && inused[i.second])
        {
            dfs_cnt(i.first, v);
            cnt[v] += cnt[i.first];
        }
}

int FindCentroid(int root)
{
    dfs_cnt(root);
    int centroid = root;

    for (pair<int, int> ans = {0, 0};;)
    {
        ans = {0, 0};
        for (auto i : adj[centroid])
            if (cnt[i.first] < cnt[centroid] && inused[i.second])
                ans = max(ans, make_pair(cnt[i.first], i.first));

        if (ans.first * 2 <= cnt[root])
            break;
        centroid = ans.second;
    }

    return centroid;
}

void dfs_fill(int v, int p = -1)
{
    su.emplace_back(du[v]);
    sd.emplace_back(dd[v]);

    for (auto i : adj[v])
        if (i.first != p && inused[i.second])
        {
            du[i.first] = du[v] - w[i.second] + a[i.first];
            dd[i.first] = dd[v] - w[i.second] + a[v];
            par[i.first] = v;
            dfs_fill(i.first, v);
        }
}

#define Unique(s) (sort(s.begin(), s.end()), s.resize(unique(s.begin(), s.end()) - s.begin()))
#define Find(s, x) (lower_bound(s.begin(), s.end(), x) - s.begin() + 1)

void CentroidDecomposition(int root)
{
    int Centroid = FindCentroid(root); // Finding Centroid

    // Compress number ??
    su.clear();
    sd.clear();
    du[Centroid] = maxn[Centroid] = 0;
    dd[Centroid] = 0;
    dfs_fill(Centroid);
    Unique(su);
    Unique(sd);
    fu.Assign(su.size());
    fd.Assign(sd.size());

    // Calculation
    fu.Update(Find(su, du[Centroid]), 1);
    fd.Update(Find(sd, dd[Centroid]), 1);

    for (auto i : adj[Centroid])
        if (inused[i.second])
        {
            queue<int> q({i.first});

            while (!q.empty())
            {
                int c = q.front();
                q.pop();

                if (maxn[par[c]] <= du[c])
                    ans += fd.Get(Find(sd, -du[c]), sd.size());

                maxn[c] = max(maxn[par[c]], du[c]);
                dd[c] = min(dd[c], dd[par[c]]);

                ans += fu.Get(Find(su, -dd[c]), su.size());

                for (auto i : adj[c])
                    if (i.first != par[c] && inused[i.second])
                        q.emplace(i.first);
            }

            q.emplace(i.first);
            while (!q.empty())
            {
                int c = q.front();
                q.pop();

                if (maxn[par[c]] <= du[c])
                    fu.Update(Find(su, du[c]), 1);
                fd.Update(Find(sd, dd[c]), 1);

                for (auto i : adj[c])
                    if (i.first != par[c] && inused[i.second])
                        q.emplace(i.first);
            }
        }

    for (auto i : adj[Centroid])
        if (inused[i.second])
        {
            inused[i.second] = 0;
            CentroidDecomposition(i.first);
        }
}

void Solve()
{
    CentroidDecomposition(1);

    cout << ans;
}

int32_t main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    if (fopen("palesta.INP", "r"))
    {
        freopen("paletsa.inp", "r", stdin);
        freopen("palesta.out", "w", stdout);
    }
    int t(1);
    if (typetest)
        cin >> t;
    for (int _ = 1; _ <= t; ++_)
    {
        // cout << "Case #" << _ << ": ";
        Read();
        Solve();
    }
    // cerr << "\nTime elapsed: " << 1000 * clock() / CLOCKS_PER_SEC << "ms\n";
}

/*
 */

Compilation message (stderr)

transport.cpp: In function 'int32_t main()':
transport.cpp:195:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  195 |         freopen("paletsa.inp", "r", stdin);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~
transport.cpp:196:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  196 |         freopen("palesta.out", "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 13 / 13
• Subtask 2: 13 / 13
• Subtask 3: 13 / 13
• Subtask 4: 13 / 13
• Subtask 5: 13 / 13
• Subtask 6: 13 / 13
• Subtask 7: 13 / 13
• Subtask 8: 13 / 13
• Subtask 9: 13 / 13
• Subtask 10: 13 / 13