Submission #857597

# Submission time Handle Problem Language Result Execution time Memory
857597 2023-10-06T12:55:01 Z boris_mihov Paths (RMI21_paths) C++17
36 / 100
600 ms 42064 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.max)
        {
            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)
    {
        if (l == r)
        {
            tree[node].max = 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);
        build(mid + 1, r, 2*node + 1);
        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 == 0)
            {
                tree[node].max = 0;
                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()
    {
        build(1, n, 1);
    }

    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 << "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)
        {
            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 == 5) std::cout << "outside max, inside min: " << outside.max << ' ' << inside.min << '\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, 0);
                treeActive.updateSet(outside.idx, outside.max);
                treeNotActive.updateSet(outside.idx, 0);

                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});
                revertActive.push_back({outside.idx, 0});
            }
        }
    }

    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)
    {
        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()
{
    for (int node = 1 ; node <= n ; ++node)
    {
        treeActive.build();
        treeNotActive.build();
        timer = 0;

        buildDFS(node, 0, 0);

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

        std::cout << treeActive.query(1, n).sum << '\n';
    }
}

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)':
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;
      |                   ~~^~~
# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35416 KB Output is correct
# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35416 KB Output is correct
3 Correct 12 ms 35420 KB Output is correct
4 Correct 12 ms 35612 KB Output is correct
5 Correct 10 ms 35416 KB Output is correct
6 Correct 12 ms 35520 KB Output is correct
7 Correct 12 ms 35420 KB Output is correct
# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35416 KB Output is correct
3 Correct 12 ms 35420 KB Output is correct
4 Correct 12 ms 35612 KB Output is correct
5 Correct 10 ms 35416 KB Output is correct
6 Correct 12 ms 35520 KB Output is correct
7 Correct 12 ms 35420 KB Output is correct
8 Correct 212 ms 35652 KB Output is correct
9 Correct 215 ms 35928 KB Output is correct
10 Correct 194 ms 35676 KB Output is correct
11 Correct 182 ms 35648 KB Output is correct
12 Correct 208 ms 35652 KB Output is correct
# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35416 KB Output is correct
3 Correct 12 ms 35420 KB Output is correct
4 Correct 12 ms 35612 KB Output is correct
5 Correct 10 ms 35416 KB Output is correct
6 Correct 12 ms 35520 KB Output is correct
7 Correct 12 ms 35420 KB Output is correct
8 Correct 212 ms 35652 KB Output is correct
9 Correct 215 ms 35928 KB Output is correct
10 Correct 194 ms 35676 KB Output is correct
11 Correct 182 ms 35648 KB Output is correct
12 Correct 208 ms 35652 KB Output is correct
13 Execution timed out 1072 ms 35924 KB Time limit exceeded
14 Halted 0 ms 0 KB -
# Verdict Execution time Memory Grader output
1 Execution timed out 1064 ms 42064 KB Time limit exceeded
2 Halted 0 ms 0 KB -
# Verdict Execution time Memory Grader output
1 Correct 6 ms 35416 KB Output is correct
2 Correct 5 ms 35416 KB Output is correct
3 Correct 12 ms 35420 KB Output is correct
4 Correct 12 ms 35612 KB Output is correct
5 Correct 10 ms 35416 KB Output is correct
6 Correct 12 ms 35520 KB Output is correct
7 Correct 12 ms 35420 KB Output is correct
8 Correct 212 ms 35652 KB Output is correct
9 Correct 215 ms 35928 KB Output is correct
10 Correct 194 ms 35676 KB Output is correct
11 Correct 182 ms 35648 KB Output is correct
12 Correct 208 ms 35652 KB Output is correct
13 Execution timed out 1072 ms 35924 KB Time limit exceeded
14 Halted 0 ms 0 KB -