oj.uz

Submission #857604

# Submission time Handle Problem Language Result Execution time Memory
857604 2023-10-06T13:29:05 Z boris_mihov Paths (RMI21_paths) C++17
100 / 100
 245 ms 48700 KB 
#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <vector>

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::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)
        {
            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 (outside.max > inside.min)
            {
                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;
    for (const auto &[u, d] : g[node])
    {
        if (u == par)
        {
            continue;
        }

        solveDFS(u, d, node);
    }

    for (const auto &[pos, val] : revertActive)
    {
        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);
    
    int used = 0;
    for (int i = 1 ; i <= k ; ++i)
    {
        SegmentTree::Node res = treeNotActive.query(1, 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);
}

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

    return 0;
}

Compilation message

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

Subtask #1
8.0 / 8.0

# Verdict Execution time Memory Grader output
1 Correct 5 ms 29020 KB Output is correct
2 Correct 4 ms 29188 KB Output is correct

Subtask #2
11.0 / 11.0

# Verdict Execution time Memory Grader output
1 Correct 5 ms 29020 KB Output is correct
2 Correct 4 ms 29188 KB Output is correct
3 Correct 5 ms 29272 KB Output is correct
4 Correct 5 ms 29112 KB Output is correct
5 Correct 5 ms 29020 KB Output is correct
6 Correct 5 ms 29020 KB Output is correct
7 Correct 5 ms 29020 KB Output is correct

Subtask #3
17.0 / 17.0

# Verdict Execution time Memory Grader output
1 Correct 5 ms 29020 KB Output is correct
2 Correct 4 ms 29188 KB Output is correct
3 Correct 5 ms 29272 KB Output is correct
4 Correct 5 ms 29112 KB Output is correct
5 Correct 5 ms 29020 KB Output is correct
6 Correct 5 ms 29020 KB Output is correct
7 Correct 5 ms 29020 KB Output is correct
8 Correct 7 ms 29020 KB Output is correct
9 Correct 6 ms 29276 KB Output is correct
10 Correct 6 ms 29144 KB Output is correct
11 Correct 7 ms 29020 KB Output is correct
12 Correct 6 ms 29276 KB Output is correct

Subtask #4
20.0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 5 ms 29020 KB Output is correct
2 Correct 4 ms 29188 KB Output is correct
3 Correct 5 ms 29272 KB Output is correct
4 Correct 5 ms 29112 KB Output is correct
5 Correct 5 ms 29020 KB Output is correct
6 Correct 5 ms 29020 KB Output is correct
7 Correct 5 ms 29020 KB Output is correct
8 Correct 7 ms 29020 KB Output is correct
9 Correct 6 ms 29276 KB Output is correct
10 Correct 6 ms 29144 KB Output is correct
11 Correct 7 ms 29020 KB Output is correct
12 Correct 6 ms 29276 KB Output is correct
13 Correct 8 ms 29272 KB Output is correct
14 Correct 8 ms 29664 KB Output is correct
15 Correct 7 ms 29276 KB Output is correct
16 Correct 8 ms 29276 KB Output is correct
17 Correct 8 ms 29276 KB Output is correct
18 Correct 6 ms 29404 KB Output is correct
19 Correct 8 ms 29300 KB Output is correct

Subtask #5
12.0 / 12.0

# Verdict Execution time Memory Grader output
1 Correct 127 ms 34588 KB Output is correct
2 Correct 160 ms 43600 KB Output is correct
3 Correct 116 ms 34248 KB Output is correct
4 Correct 131 ms 34548 KB Output is correct
5 Correct 139 ms 37680 KB Output is correct
6 Correct 124 ms 34508 KB Output is correct

Subtask #6
32.0 / 32.0

# Verdict Execution time Memory Grader output
1 Correct 5 ms 29020 KB Output is correct
2 Correct 4 ms 29188 KB Output is correct
3 Correct 5 ms 29272 KB Output is correct
4 Correct 5 ms 29112 KB Output is correct
5 Correct 5 ms 29020 KB Output is correct
6 Correct 5 ms 29020 KB Output is correct
7 Correct 5 ms 29020 KB Output is correct
8 Correct 7 ms 29020 KB Output is correct
9 Correct 6 ms 29276 KB Output is correct
10 Correct 6 ms 29144 KB Output is correct
11 Correct 7 ms 29020 KB Output is correct
12 Correct 6 ms 29276 KB Output is correct
13 Correct 8 ms 29272 KB Output is correct
14 Correct 8 ms 29664 KB Output is correct
15 Correct 7 ms 29276 KB Output is correct
16 Correct 8 ms 29276 KB Output is correct
17 Correct 8 ms 29276 KB Output is correct
18 Correct 6 ms 29404 KB Output is correct
19 Correct 8 ms 29300 KB Output is correct
20 Correct 127 ms 34588 KB Output is correct
21 Correct 160 ms 43600 KB Output is correct
22 Correct 116 ms 34248 KB Output is correct
23 Correct 131 ms 34548 KB Output is correct
24 Correct 139 ms 37680 KB Output is correct
25 Correct 124 ms 34508 KB Output is correct
26 Correct 232 ms 34924 KB Output is correct
27 Correct 187 ms 43512 KB Output is correct
28 Correct 210 ms 45392 KB Output is correct
29 Correct 116 ms 34476 KB Output is correct
30 Correct 210 ms 35024 KB Output is correct
31 Correct 165 ms 34640 KB Output is correct
32 Correct 245 ms 38740 KB Output is correct
33 Correct 211 ms 34840 KB Output is correct
34 Correct 119 ms 34808 KB Output is correct
35 Correct 222 ms 34804 KB Output is correct
36 Correct 234 ms 48700 KB Output is correct