Submission #310536

#	Time	Username	Problem	Language	Result	Execution time	Memory
310536		phathnv	Transport (COCI19_transport)	C++11	0 / 130	3 ms	2720 KiB

Main

#include <bits/stdc++.h>

#define mp make_pair
#define X first
#define Y second
#define taskname "Transport"

using namespace std;

typedef long long ll;
typedef pair <int, int> ii;

const int N = 1e5 + 1;

struct Tbit{
    int d[N];
    void update(int x, int val){
        for(; x < N; x += x & -x)
            d[x] += val;
    }
    int get(int x){
        int res = 0;
        for(; x > 0; x -= x & -x)
            res += d[x];
        return res;
    }
} bit;

int n, a[N];
vector <ii> adj[N];

bool done[N];
int sz[N];
ll sumA[N], sumW[N], down[N], up[N];
vector <int> vst;
vector <ll> vals;

void readInput(){
    scanf("%d", &n);
    for(int i = 1; i <= n; i++)
        scanf("%d", &a[i]);
    for(int i = 1; i < n; i++){
        int u, v, w;
        scanf("%d %d %d", &u, &v, &w);
        adj[u].push_back(mp(v, w));
        adj[v].push_back(mp(u, w));
    }
}

void dfsForCentroid(int u, int p = 0){
    sz[u] = 1;
    for(ii e : adj[u]){
        int v = e.X;
        if (v == p || done[v])
            continue;
        dfsForCentroid(v, u);
        sz[u] += sz[v];
    }
}

int findCentroid(int u, int p, int root){
    for(ii e : adj[u]){
        int v = e.X;
        if (v == p || done[v])
            continue;
        if (sz[v] >= sz[root] / 2)
            return findCentroid(v, u, root);
    }
    return u;
}

void dfs(int u, int p = 0){
    vst.push_back(u);
    for(ii e : adj[u]){
        int v = e.X;
        int w = e.Y;
        if (v == p || done[v])
            continue;
        sumA[v] = sumA[u] + a[v];
        sumW[v] = sumW[u] + w;
        down[v] = max(down[u], sumW[v] - sumA[u]);
        up[v] = max((ll) w - a[v], up[u] + w - a[v]);
        dfs(v, u);
    }
}

void upd(int u, int p, int type){
    int ind = lower_bound(vals.begin(), vals.end(), down[u]) - vals.begin() + 1;
    bit.update(ind, type);
    for(ii e : adj[u]){
        int v = e.X;
        if (v == p || done[v])
            continue;
        upd(v, u, type);
    }
}

int get(int u, int p = 0){
    int res = 0;
    if (up[u] <= 0){
        ll val = sumA[u] - sumW[u];
        int ind = upper_bound(vals.begin(), vals.end(), val) - vals.begin();
        res += bit.get(ind);
    }
    for(ii e : adj[u]){
        int v = e.X;
        if (v == p || done[v])
            continue;
        res += get(v, u);
    }
    return res;
}

ll calc(int u){
    dfsForCentroid(u);
    u = findCentroid(u, 0, u);

    vst.clear();
    vals.clear();
    sumW[u] = 0;
    sumA[u] = a[u];
    up[u] = down[u] = 0;
    dfs(u);
    for(int v : vst){
        vals.push_back(down[v]);
        sumA[v] -= a[u];
    }
    sort(vals.begin(), vals.end());

    ll res = 0;

    for(int trial = 0; trial < 2; trial++){
        if (trial == 0){
            res += bit.get(1);
            bit.update(1, 1);
        }
        for(ii e : adj[u]){
            int v = e.X;
            if (done[v])
                continue;
            res += get(v, u);
            upd(v, u, 1);
        }
        if (trial == 1){
            res += bit.get(1);
            bit.update(1, 1);
        }

        for(ii e : adj[u]){
            int v = e.X;
            if (done[v])
                continue;
            upd(v, u, -1);
        }

        bit.update(1, -1);
        reverse(adj[u].begin(), adj[u].end());
    }

    done[u] = 1;
    for(ii e : adj[u]){
        int v = e.X;
        if (!done[v])
            res += calc(v);
    }
    return res;
}

void solve(){
    cout << calc(1);
}

int main(){
    freopen(taskname".inp", "r", stdin);
    freopen(taskname".out", "w", stdout);
    readInput();
    solve();
    return 0;
}

Compilation message (stderr)

transport.cpp: In function 'void readInput()':
transport.cpp:39:10: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   39 |     scanf("%d", &n);
      |     ~~~~~^~~~~~~~~~
transport.cpp:41:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   41 |         scanf("%d", &a[i]);
      |         ~~~~~^~~~~~~~~~~~~
transport.cpp:44:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   44 |         scanf("%d %d %d", &u, &v, &w);
      |         ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
transport.cpp: In function 'int main()':
transport.cpp:174:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
  174 |     freopen(taskname".inp", "r", stdin);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~
transport.cpp:175:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
  175 |     freopen(taskname".out", "w", stdout);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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