oj.uz

Submission #857602

# Submission time Handle Problem Language Result Execution time Memory
857602 2023-10-06T13:25:16 Z boris_mihov Paths (RMI21_paths) C++17
100 / 100
 292 ms 67724 KB 
#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

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.0 / 8.0

# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35672 KB Output is correct

Subtask #2
11.0 / 11.0

# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35672 KB Output is correct
3 Correct 6 ms 35420 KB Output is correct
4 Correct 6 ms 35564 KB Output is correct
5 Correct 6 ms 35504 KB Output is correct
6 Correct 6 ms 35420 KB Output is correct
7 Correct 5 ms 35420 KB Output is correct

Subtask #3
17.0 / 17.0

# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35672 KB Output is correct
3 Correct 6 ms 35420 KB Output is correct
4 Correct 6 ms 35564 KB Output is correct
5 Correct 6 ms 35504 KB Output is correct
6 Correct 6 ms 35420 KB Output is correct
7 Correct 5 ms 35420 KB Output is correct
8 Correct 7 ms 35416 KB Output is correct
9 Correct 7 ms 35932 KB Output is correct
10 Correct 7 ms 35500 KB Output is correct
11 Correct 6 ms 35420 KB Output is correct
12 Correct 7 ms 35528 KB Output is correct

Subtask #4
20.0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35672 KB Output is correct
3 Correct 6 ms 35420 KB Output is correct
4 Correct 6 ms 35564 KB Output is correct
5 Correct 6 ms 35504 KB Output is correct
6 Correct 6 ms 35420 KB Output is correct
7 Correct 5 ms 35420 KB Output is correct
8 Correct 7 ms 35416 KB Output is correct
9 Correct 7 ms 35932 KB Output is correct
10 Correct 7 ms 35500 KB Output is correct
11 Correct 6 ms 35420 KB Output is correct
12 Correct 7 ms 35528 KB Output is correct
13 Correct 9 ms 35672 KB Output is correct
14 Correct 9 ms 36444 KB Output is correct
15 Correct 8 ms 35932 KB Output is correct
16 Correct 8 ms 35676 KB Output is correct
17 Correct 9 ms 35800 KB Output is correct
18 Correct 8 ms 35676 KB Output is correct
19 Correct 10 ms 35672 KB Output is correct

Subtask #5
12.0 / 12.0

# Verdict Execution time Memory Grader output
1 Correct 159 ms 43436 KB Output is correct
2 Correct 186 ms 59188 KB Output is correct
3 Correct 123 ms 45704 KB Output is correct
4 Correct 137 ms 45480 KB Output is correct
5 Correct 167 ms 50276 KB Output is correct
6 Correct 138 ms 45256 KB Output is correct

Subtask #6
32.0 / 32.0

# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35672 KB Output is correct
3 Correct 6 ms 35420 KB Output is correct
4 Correct 6 ms 35564 KB Output is correct
5 Correct 6 ms 35504 KB Output is correct
6 Correct 6 ms 35420 KB Output is correct
7 Correct 5 ms 35420 KB Output is correct
8 Correct 7 ms 35416 KB Output is correct
9 Correct 7 ms 35932 KB Output is correct
10 Correct 7 ms 35500 KB Output is correct
11 Correct 6 ms 35420 KB Output is correct
12 Correct 7 ms 35528 KB Output is correct
13 Correct 9 ms 35672 KB Output is correct
14 Correct 9 ms 36444 KB Output is correct
15 Correct 8 ms 35932 KB Output is correct
16 Correct 8 ms 35676 KB Output is correct
17 Correct 9 ms 35800 KB Output is correct
18 Correct 8 ms 35676 KB Output is correct
19 Correct 10 ms 35672 KB Output is correct
20 Correct 159 ms 43436 KB Output is correct
21 Correct 186 ms 59188 KB Output is correct
22 Correct 123 ms 45704 KB Output is correct
23 Correct 137 ms 45480 KB Output is correct
24 Correct 167 ms 50276 KB Output is correct
25 Correct 138 ms 45256 KB Output is correct
26 Correct 259 ms 45864 KB Output is correct
27 Correct 208 ms 59204 KB Output is correct
28 Correct 250 ms 62392 KB Output is correct
29 Correct 136 ms 45692 KB Output is correct
30 Correct 245 ms 45768 KB Output is correct
31 Correct 182 ms 45292 KB Output is correct
32 Correct 292 ms 51796 KB Output is correct
33 Correct 189 ms 45908 KB Output is correct
34 Correct 133 ms 46160 KB Output is correct
35 Correct 224 ms 45864 KB Output is correct
36 Correct 260 ms 67724 KB Output is correct