Submission #1046297

#	Time	Username	Problem	Language	Result	Execution time	Memory
1046297		SamAnd	Petrol stations (CEOI24_stations)	C++17	100 / 100	798 ms	33984 KiB

Main

#include <bits/stdc++.h>
using namespace std;
#define m_p make_pair
#define all(x) (x).begin(),(x).end()
#define sz(x) ((int)(x).size())
#define fi first
#define se second
typedef long long ll;
mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());
mt19937 rnf(2106);
const int N = 70004;

int n;
ll k;
vector<pair<int, ll> > g[N];

bool c[N];
int q[N];
void dfs0(int x, int p)
{
    q[x] = 1;
    for (int i = 0; i < g[x].size(); ++i)
    {
        int h = g[x][i].fi;
        if (h == p)
            continue;
        if (c[h])
            continue;
        dfs0(h, x);
        q[x] += q[h];
    }
}

int dfsc(int x, int p, int xx)
{
    for (int i = 0; i < g[x].size(); ++i)
    {
        int h = g[x][i].fi;
        if (h == p)
            continue;
        if (c[h])
            continue;
        if (q[h] > q[xx] / 2)
            return dfsc(h, x, xx);
    }
    return x;
}

ll ans[N];

const int m = 17;
int pp[N][m];
ll ss[N][m];
int tin[N], tout[N], ti;
void dfs1(int x, int p, ll dp)
{
    tin[x] = ++ti;
    pp[x][0] = p;
    ss[x][0] = dp;
    for (int i = 1; i < m; ++i)
    {
        pp[x][i] = pp[pp[x][i - 1]][i - 1];
        ss[x][i] = ss[x][i - 1] + ss[pp[x][i - 1]][i - 1];
    }
    for (int i = 0; i < g[x].size(); ++i)
    {
        int h = g[x][i].fi;
        ll d = g[x][i].se;
        if (h == p)
            continue;
        if (c[h])
            continue;
        dfs1(h, x, d);
    }
    tout[x] = ti;
}

bool isp(int x, int y)
{
    return (tin[x] <= tin[y] && tin[y] <= tout[x]);
}

int FU(int x, ll k)
{
    ll s = 0;
    for (int i = m - 1; i >= 0; --i)
    {
        if (s + ss[x][i] <= k)
        {
            s += ss[x][i];
            x = pp[x][i];
        }
    }
    return x;
}

void dfst(int x, int p, vector<int>& v)
{
    for (int i = 0; i < g[x].size(); ++i)
    {
        int h = g[x][i].fi;
        if (h == p)
            continue;
        if (c[h])
            continue;
        dfst(h, x, v);
    }
    v.push_back(x);
}

void ubd(vector<int>& t, int tl, int tr, int x, int y, int pos)
{
    assert(tl <= x && x <= tr);
    if (tl == tr)
    {
        t[pos] += y;
        return;
    }
    int m = (tl + tr) / 2;
    if (x <= m)
        ubd(t, tl, m, x, y, pos * 2);
    else
        ubd(t, m + 1, tr, x, y, pos * 2 + 1);
    t[pos] = t[pos * 2] + t[pos * 2 + 1];
}

int qry(vector<int>& t, int tl, int tr, int l, int r, int pos)
{
    if (l > r)
        return 0;
    assert(tl <= l && r <= tr);
    if (tl == l && tr == r)
        return t[pos];
    int m = (tl + tr) / 2;
    return (qry(t, tl, m, l, min(m, r), pos * 2) +
            qry(t, m + 1, tr, max(m + 1, l), r, pos * 2 + 1));
}

void dfsf(int x, int p, const vector<ll>& xt, vector<int>& t, int root, vector<int>& hav)
{
    int v = 0;
    ll r = k - ss[pp[x][0]][m - 1];
    ll l = k - ss[x][m - 1] + 1;
    if (l <= r)
    {
        int l2 = lower_bound(all(xt), l) - xt.begin() + 1;
        int r2 = upper_bound(all(xt), r) - xt.begin();
        if (l2 <= r2)
        {
            assert(1 <= l2 && r2 <= sz(xt));
            v += qry(t, 1, sz(xt), l2, r2, 1);
        }
    }
    if (ss[x][m - 1] > k && ss[pp[x][0]][m - 1] <= k)
    {
        ++v;
    }
    if (ss[x][m - 1] > k)
    {
        int y = FU(pp[x][0], k);
        int z = pp[FU(x, k)][0];
        v += qry(hav, 1, q[root], tin[y], tin[z], 1);
    }

    ans[pp[x][0]] += v * 1LL * q[x];
    ubd(hav, 1, q[root], tin[pp[x][0]], v, 1);
    for (int i = 0; i < g[x].size(); ++i)
    {
        int h = g[x][i].fi;
        if (h == p)
            continue;
        if (c[h])
            continue;
        dfsf(h, x, xt, t, root, hav);
    }

    ubd(hav, 1, q[root], tin[pp[x][0]], -v, 1);
}

int havv[N];
void solvv(int r)
{
    dfs0(r, r);
    ti = 0;
    dfs1(r, r, 0);
    map<int, vector<pair<ll, int> > > mp;
    for (int i = 0; i < g[r].size(); ++i)
    {
        int x = g[r][i].fi;
        int cx = x;
        if (c[x])
            continue;
        vector<int> v;
        dfst(x, r, v);
        for (int i = 0; i < sz(v); ++i)
            havv[v[i]] = 1;
        for (int i = 0; i < sz(v); ++i)
        {
            int x = v[i];
            int y = FU(x, k);
            if (y == r)
            {
                mp[cx].push_back(m_p(ss[x][m - 1], havv[x]));
            }
            else
            {
                ans[y] += havv[x] * 1LL * (q[r] - sz(v));
                havv[y] += havv[x];
            }
        }
    }

    vector<ll> xt;
    for (auto it = mp.begin(); it != mp.end(); ++it)
    {
        for (int i = 0; i < sz(it->se); ++i)
        {
            xt.push_back(it->se[i].fi);
        }
    }
    sort(all(xt));

    vector<int> t(sz(xt) * 4 + 10, 0);
    for (auto it = mp.begin(); it != mp.end(); ++it)
    {
        for (int i = 0; i < sz(it->se); ++i)
        {
            ubd(t, 1, sz(xt), lower_bound(all(xt), it->se[i].fi) - xt.begin() + 1, it->se[i].se, 1);
        }
    }

    vector<int> hav(q[r] * 4 + 10, 0);
    for (int i = 0; i < g[r].size(); ++i)
    {
        int x = g[r][i].fi;
        if (c[x])
            continue;
        vector<int> v;
        dfst(x, r, v);
        reverse(all(v));
        for (int i = 0; i < sz(mp[x]); ++i)
        {
            ubd(t, 1, sz(xt), lower_bound(all(xt), mp[x][i].fi) - xt.begin() + 1, -mp[x][i].se, 1);
        }
        dfsf(x, r, xt, t, r, hav);
        for (int i = 0; i < sz(mp[x]); ++i)
        {
            ubd(t, 1, sz(xt), lower_bound(all(xt), mp[x][i].fi) - xt.begin() + 1, mp[x][i].se, 1);
        }
    }
}

vector<ll> solv()
{
    for (int i = 0; i <= n + 1; ++i)
    {
        c[i] = false;
        ans[i] = 0;
    }

    queue<int> q;
    q.push(1);
    while (!q.empty())
    {
        int x = q.front();
        q.pop();
        dfs0(x, x);
        x = dfsc(x, x, x);

        solvv(x);

        c[x] = true;
        for (int i = 0; i < g[x].size(); ++i)
        {
            int h = g[x][i].fi;
            if (!c[h])
                q.push(h);
        }
    }

    vector<ll> ansv;
    for (int x = 1; x <= n; ++x)
        ansv.push_back(ans[x]);
    for (int x = 1; x <= n; ++x)
        cout << ans[x] << "\n";
    return ansv;
}

bool dfsg(int x, int p, int y, vector<ll>& v, vector<int>& u)
{
    if (x == y)
        return true;
    for (int i = 0; i < g[x].size(); ++i)
    {
        int h = g[x][i].fi;
        if (h == p)
            continue;
        v.push_back(g[x][i].se);
        u.push_back(x);
        if (dfsg(h, x, y, v, u))
            return true;
        u.pop_back();
        v.pop_back();
    }
    return false;
}

vector<ll> solv0()
{
    for (int i = 0; i <= n + 1; ++i)
        ans[i] = 0;
    for (int x = 1; x <= n; ++x)
    {
        for (int y = 1; y <= n; ++y)
        {
            vector<ll> v;
            vector<int> u;
            assert(dfsg(x, x, y, v, u));
            ll s = 0;
            for (int i = 0; i < sz(v); ++i)
            {
                if (s + v[i] > k)
                {
                    s = 0;
                    ++ans[u[i]];
                }
                s += v[i];
            }
        }
    }
    vector<ll> ansv;
    for (int x = 1; x <= n; ++x)
        ansv.push_back(ans[x]);
    return ansv;
}

int main()
{
    #ifdef SOMETHING
    freopen("input.txt", "r", stdin);
    //freopen("output.txt", "w", stdout);
    #endif // SOMETHING
    ios_base::sync_with_stdio(false), cin.tie(0);

    /*int st = 1000;
    while (st--)
    {
        n = rnf() % 10 + 2;
        k = rnf() % 5 + 1;
        for (int x = 1; x <= n; ++x)
            g[x].clear();
        vector<pair<pair<int, int>, int> > b;
        for (int x = 2; x <= n; ++x)
        {
            int y = rnf() % (x - 1) + 1;
            int z = rnf() % (k + 1);
            g[x].push_back(m_p(y, z));
            g[y].push_back(m_p(x, z));
            b.push_back(m_p(m_p(x, y), z));
        }
        auto myans = solv();
        auto trueans = solv0();
        if (myans != trueans)
        {
            solv();
            solv0();
            cout << "WA\n";
            cout << n << ' ' << k << "\n";
            for (int i = 0; i < n - 1; ++i)
                cout << b[i].fi.fi << ' ' << b[i].fi.se << ' ' << b[i].se << "\n";
            cout << "MY\n";
            for (int i = 0; i < sz(myans); ++i)
                cout << myans[i] << "\n";
            cout << "TRUE\n";
            for (int i = 0; i < sz(trueans); ++i)
                cout << trueans[i] << "\n";
            cout << endl;
            return 0;
        }
    }
    cout << "OK" << endl;
    return 0;*/

    int tt = 1;
    //cin >> tt;
    while (tt--)
    {
        cin >> n >> k;
        for (int i = 0; i < n - 1; ++i)
        {
            int x, y, z;
            cin >> x >> y >> z;
            ++x;
            ++y;
            g[x].push_back(m_p(y, z));
            g[y].push_back(m_p(x, z));
        }
        solv();
    }
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void dfs0(int, int)':
Main.cpp:22:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   22 |     for (int i = 0; i < g[x].size(); ++i)
      |                     ~~^~~~~~~~~~~~~
Main.cpp: In function 'int dfsc(int, int, int)':
Main.cpp:36:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   36 |     for (int i = 0; i < g[x].size(); ++i)
      |                     ~~^~~~~~~~~~~~~
Main.cpp: In function 'void dfs1(int, int, ll)':
Main.cpp:65:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   65 |     for (int i = 0; i < g[x].size(); ++i)
      |                     ~~^~~~~~~~~~~~~
Main.cpp: In function 'void dfst(int, int, std::vector<int>&)':
Main.cpp:99:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   99 |     for (int i = 0; i < g[x].size(); ++i)
      |                     ~~^~~~~~~~~~~~~
Main.cpp: In function 'void dfsf(int, int, const std::vector<long long int>&, std::vector<int>&, int, std::vector<int>&)':
Main.cpp:167:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  167 |     for (int i = 0; i < g[x].size(); ++i)
      |                     ~~^~~~~~~~~~~~~
Main.cpp: In function 'void solvv(int)':
Main.cpp:187:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  187 |     for (int i = 0; i < g[r].size(); ++i)
      |                     ~~^~~~~~~~~~~~~
Main.cpp:233:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  233 |     for (int i = 0; i < g[r].size(); ++i)
      |                     ~~^~~~~~~~~~~~~
Main.cpp: In function 'std::vector<long long int> solv()':
Main.cpp:273:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  273 |         for (int i = 0; i < g[x].size(); ++i)
      |                         ~~^~~~~~~~~~~~~
Main.cpp: In function 'bool dfsg(int, int, int, std::vector<long long int>&, std::vector<int>&)':
Main.cpp:293:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  293 |     for (int i = 0; i < g[x].size(); ++i)
      |                     ~~^~~~~~~~~~~~~

• Subtask 0: 0 / 0
• Subtask 1: 18 / 18
• Subtask 2: 8 / 8
• Subtask 3: 10 / 10
• Subtask 4: 12 / 12
• Subtask 5: 17 / 17
• Subtask 6: 35 / 35