Submission #857602

#	Time	Username	Problem	Language	Result	Execution time	Memory
857602		boris_mihov	Paths (RMI21_paths)	C++17	100 / 100	292 ms	67724 KiB

Main

#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <vector>

#define int long long
typedef long long llong;
const int MAXN = 100000 + 10;
const llong INF = 1e18;
const int INTINF = 1e9;

int n, k;
struct SegmentTree
{
    struct Node
    {
        llong sum;
        llong max;
        llong min;
        int idxMin;
        int idx;

        Node()
        {
            max = -1;
            min = INF;
            idxMin = 0;
            idx = 0;
            sum = 0;
        }
    };

    Node tree[4*MAXN];
    Node combine(Node left, Node right)
    {
        Node res;
        res.sum = left.sum + right.sum;
        if (left.max > right.max)
        {
            res.max = left.max;
            res.idx = left.idx;
        } else
        {
            res.max = right.max;
            res.idx = right.idx;
        }

        if (left.min < right.min)
        {
            res.min = left.min;
            res.idxMin = left.idxMin;
        } else
        {
            res.min = right.min;
            res.idxMin = right.idxMin;
        }

        return res;
    }

    void build(int l, int r, int node, bool inactive)
    {
        if (l == r)
        {
            tree[node].max = (inactive ? -1 : 0);
            tree[node].sum = 0;
            tree[node].min = INF;
            tree[node].idx = l;
            tree[node].idxMin = l;
            return;
        }

        int mid = l + r >> 1;
        build(l, mid, 2*node, inactive);
        build(mid + 1, r, 2*node + 1, inactive);
        tree[node] = combine(tree[2*node], tree[2*node + 1]);
    }

    void updateSet(int l, int r, int node, int queryPos, llong queryValue)
    {
        if (l == r)
        {
            if (queryValue == -1)
            {
                tree[node].max = -1;
                tree[node].sum = 0;
                tree[node].min = INF;
            } else
            {
                tree[node].max = queryValue;
                tree[node].min = queryValue;
                tree[node].sum = queryValue;
            }

            return;
        }

        int mid = l + r >> 1;
        if (queryPos <= mid) updateSet(l, mid, 2*node, queryPos, queryValue);
        else updateSet(mid + 1, r, 2*node + 1, queryPos, queryValue);
        tree[node] = combine(tree[2*node], tree[2*node + 1]);
    }

    void updateAdd(int l, int r, int node, int queryPos, llong queryValue)
    {
        if (l == r)
        {
            tree[node].max += queryValue;
            tree[node].min += queryValue;
            tree[node].sum += queryValue;
            return;
        }

        int mid = l + r >> 1;
        if (queryPos <= mid) updateAdd(l, mid, 2*node, queryPos, queryValue);
        else updateAdd(mid + 1, r, 2*node + 1, queryPos, queryValue);
        tree[node] = combine(tree[2*node], tree[2*node + 1]);
    }
    
    Node query(int l, int r, int node, int queryL, int queryR)
    {
        if (queryL <= l && r <= queryR)
        {
            return tree[node];
        }

        int mid = l + r >> 1;
        Node res;

        if (queryL <= mid) res = combine(res, query(l, mid, 2*node, queryL, queryR));
        if (mid + 1 <= queryR) res = combine(res, query(mid + 1, r, 2*node + 1, queryL, queryR));
        return res;
    }

    void build(bool inactive)
    {
        build(1, n, 1, inactive);
    }

    void updateSet(int pos, llong value)
    {
        updateSet(1, n, 1, pos, value);
    }

    void updateAdd(int pos, llong value)
    {
        updateAdd(1, n, 1, pos, value);
    }

    Node query(int l, int r)
    {
        return query(1, n, 1, l, r);
    }
};

struct Fenwick
{
    int tree[MAXN];
    void update(int pos, int value)
    {
        for (int idx = pos ; idx <= n ; idx += idx & (-idx))
        {
            tree[idx] += value;
        }
    }

    int query(int pos)
    {
        int res = 0;
        for (int idx = pos ; idx > 0 ; idx -= idx & (-idx))
        {
            res += tree[idx];
        }

        return res;
    }

    int query(int l, int r)
    {
        return query(r) - query(l - 1);
    }
};

int timer;
int in[MAXN];
int out[MAXN];
SegmentTree treeActive;
SegmentTree treeNotActive;
std::vector <std::pair <int,int>> g[MAXN];
Fenwick fenwick;
llong ans[MAXN];

void buildDFS(int node, int inc, int par)
{
    in[node] = ++timer;
    for (const auto &[u, d] : g[node])
    {
        if (u == par)
        {
            continue;
        }

        buildDFS(u, d, node);
    }

    out[node] = timer;
    SegmentTree::Node res = treeNotActive.query(in[node], out[node]);
    treeNotActive.updateAdd(res.idx, inc);
};

void solveDFS(int node, int inc, int par)
{
    // std::cout << "trees before: " << node << '\n';
    // for (int i = 1 ; i <= n ; ++i)
    // {
    //     std::cout << i << ": " << treeActive.query(in[i], in[i]).max << ' ' << treeActive.query(in[i], in[i]).min << ' ' << treeNotActive.query(in[i], in[i]).max << ' ' << fenwick.query(in[i], in[i]) << '\n';
    // }

    // std::cout << "solve dfs: " << node << ' ' << inc << ' ' << treeActive.query(1, n).sum << '\n' << std::flush;
    std::vector <std::pair <int,llong>> revertActive;    
    std::vector <std::pair <int,llong>> revertNotActive;    
    std::vector <std::pair <int,int>> revertFenwick;    

    if (inc != 0)
    {
        if (fenwick.query(1, n) - fenwick.query(in[node], out[node]) != 0)
        {
            SegmentTree::Node outside;
            if (in[node] > 1)
            {
                outside = treeActive.combine(outside, treeActive.query(1, in[node] - 1));
            }

            if (out[node] < n)
            {
                outside = treeActive.combine(outside, treeActive.query(out[node] + 1, n));
            }

            treeActive.updateAdd(outside.idx, inc);
            revertActive.push_back({outside.idx, outside.max});
        } else
        {
            SegmentTree::Node outside;
            if (in[node] > 1)
            {
                outside = treeActive.combine(outside, treeNotActive.query(1, in[node] - 1));
            }

            if (out[node] < n)
            {
                outside = treeActive.combine(outside, treeNotActive.query(out[node] + 1, n));
            }

            treeNotActive.updateAdd(outside.idx, inc);
            revertNotActive.push_back({outside.idx, outside.max});
        }

        if (fenwick.query(in[node], out[node]) != 0)
        {
            // std::cout << "here\n";
            SegmentTree::Node inside = treeActive.query(in[node], out[node]);
            int removedIdx = inside.idx;

            treeActive.updateAdd(inside.idx, -inc);
            revertActive.push_back({inside.idx, inside.max});
            inside = treeActive.query(in[node], out[node]);

            SegmentTree::Node outside;
            if (in[node] > 1)
            {
                outside = treeActive.combine(outside, treeNotActive.query(1, in[node] - 1));
            }

            if (out[node] < n)
            {
                outside = treeActive.combine(outside, treeNotActive.query(out[node] + 1, n));
            }

            // if (node == 2) 
            // {
            //     std::cout << "outside max, inside min: " << outside.max << ' ' << inside.min << ' ' << in[node] << ' ' << out[node] << '\n';
            //     for (int i = 1 ; i <= n ; ++i)
            //     {
            //         std::cout << i << ": " << treeActive.query(in[i], in[i]).max << ' ' << treeActive.query(in[i], in[i]).min << ' ' << treeNotActive.query(in[i], in[i]).max << ' ' << fenwick.query(in[i], in[i]) << '\n';
            //     }
            // }

            if (outside.max > inside.min)
            {
                // std::cout << "max min: " << outside.max << ' ' << inside.min << ' ' << inside.idxMin << '\n' << std::flush;
                revertFenwick.push_back({inside.idxMin, 1});
                revertFenwick.push_back({outside.idx, -1});
                fenwick.update(inside.idxMin, -1);
                fenwick.update(outside.idx, 1);

                treeActive.updateSet(inside.idxMin, -1);
                treeNotActive.updateSet(inside.idxMin, inside.min);
                treeActive.updateSet(outside.idx, outside.max);
                treeNotActive.updateSet(outside.idx, -1);

                if (removedIdx == inside.idxMin) revertActive.push_back({inside.idxMin, inside.min + inc});
                else revertActive.push_back({inside.idxMin, inside.min});
                revertNotActive.push_back({outside.idx, outside.max});
                revertNotActive.push_back({inside.idxMin, -1});
                revertActive.push_back({outside.idx, -1});
            }
        }
    }

    ans[node] = treeActive.query(1, n).sum;
    // std::cout << "ans: " << node << " = " << ans[node] << '\n' << std::flush;
    for (const auto &[u, d] : g[node])
    {
        if (u == par)
        {
            continue;
        }

        solveDFS(u, d, node);
        // if (node == 1)
        // {
        //     std::cout << "trees after: " << u << '\n';
        //     for (int i = 1 ; i <= n ; ++i)
        //     {
        //         std::cout << i << ": " << treeActive.query(in[i], in[i]).max << ' ' << treeNotActive.query(in[i], in[i]).max << ' ' << fenwick.query(in[i], in[i]) << '\n';
        //     }
        // }

        assert(treeActive.query(1, n).sum == ans[node]);
    }

    for (const auto &[pos, val] : revertActive)
    {
        // std::cout << "revert active: " << pos << ' ' << val << '\n';
        treeActive.updateSet(pos, val);
    }

    for (const auto &[pos, val] : revertNotActive)
    {
        treeNotActive.updateSet(pos, val);
    }
    
    for (const auto &[pos, val] : revertFenwick)
    {
        fenwick.update(pos, val);
    }
}

void solve()
{
    treeActive.build(true);
    treeNotActive.build(false);
    buildDFS(1, 0, 0);

    // std::cout << "build\n";
    // for (int i = 1 ; i <= n ; ++i)
    // {
    //     std::cout << i << ": " << treeNotActive.query(in[i], in[i]).max << '\n';
    // }
    
    int used = 0;
    for (int i = 1 ; i <= k ; ++i)
    {
        SegmentTree::Node res = treeNotActive.query(1, n);
        // std::cout << "active: " << res.idx << ' ' << res.max << '\n';
        treeNotActive.updateSet(res.idx, -1);
        treeActive.updateSet(res.idx, res.max);
        fenwick.update(res.idx, 1);
    }

    solveDFS(1, 0, 0);
}

void print()
{
    for (int i = 1 ; i <= n ; ++i)
    {
        std::cout << ans[i] << '\n';
    }
}

void input()
{
    std::cin >> n >> k;
    for (int i = 1 ; i < n ; ++i)
    {
        int u, v, d;
        std::cin >> u >> v >> d;
        g[u].push_back({v, d});
        g[v].push_back({u, d});
    }
}

void fastIOI()
{
    std::ios_base :: sync_with_stdio(0);
    std::cout.tie(nullptr);
    std::cin.tie(nullptr);
}

signed main()
{
    fastIOI();
    input();
    solve();
    print();

    return 0;
}

Compilation message (stderr)

Main.cpp: In member function 'void SegmentTree::build(long long int, long long int, long long int, bool)':
Main.cpp:74:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   74 |         int mid = l + r >> 1;
      |                   ~~^~~
Main.cpp: In member function 'void SegmentTree::updateSet(long long int, long long int, long long int, long long int, llong)':
Main.cpp:99:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   99 |         int mid = l + r >> 1;
      |                   ~~^~~
Main.cpp: In member function 'void SegmentTree::updateAdd(long long int, long long int, long long int, long long int, llong)':
Main.cpp:115:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
  115 |         int mid = l + r >> 1;
      |                   ~~^~~
Main.cpp: In member function 'SegmentTree::Node SegmentTree::query(long long int, long long int, long long int, long long int, long long int)':
Main.cpp:128:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
  128 |         int mid = l + r >> 1;
      |                   ~~^~~
Main.cpp: In function 'void solve()':
Main.cpp:362:9: warning: unused variable 'used' [-Wunused-variable]
  362 |     int used = 0;
      |         ^~~~

• Subtask 1: 8 / 8
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 20 / 20
• Subtask 5: 12 / 12
• Subtask 6: 32 / 32