Submission #1241017

#	Time	Username	Problem	Language	Result	Execution time	Memory
1241017		ProtonDecay314	Sprinkler (JOI22_sprinkler)	C++20	0 / 100	1346 ms	70836 KiB

sprinkler

/*
stuff i need to do:
- might need to make a permutation inverter
- compute preorder, make a translation table from tree index to preorder index
- compute subtree sizes; make sure to index by preorder
- compute BFS order, make a translation table from tree index to BFS index (and vice versa)
- segtree over BFS order

subroutine: binary search for the first node that's a subtree of some other node and is at some given depth
observation: dfs index is monotonic over a single level of a tree
first node invariant: l = set of all nodes such that dep[i] < dep_targ or dep[i] == dep_targ and dfs_order[i] < dfs_order[targ]
r is the first node

last node invariant: r = set of all nodes such that dep[i] > dep_targ or dep[i] == dep_targ and not within tree and sz[dfs_order[targ]] < dfs_order[i] - dfs_order[targ] + 1ll
*/
#pragma GCC optimize("O3", "unroll-loops")
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
typedef vector<vll> vvll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<vvi> v3i;
typedef vector<v3i> v4i;
typedef vector<bool> vb;
typedef vector<vb> vvb;
typedef pair<int, int> pi;
typedef pair<ll, ll> pll;
typedef vector<pi> vpi;
typedef vector<pll> vpll;
#define INF(dt) numeric_limits<dt>::max()
#define NINF(dt) numeric_limits<dt>::min()
#define pb push_back

ll MOD = 1ll; // TODO set this to the correct value later on!

struct Tree {
    Tree *lt, *rt;
    ll l, r, v;
    ll lazy;
    bool marked;

    Tree(ll a_l, ll a_r): lt(nullptr), rt(nullptr), l(a_l), r(a_r), lazy(1ll), marked(false) {};

    void push() {
        if(!marked) return;

        if(lt == nullptr) {
            marked = false;
            lazy = 1ll;
            return;
        }

        lt->lazy = (lt->lazy * lazy) % MOD;
        lt->v = (lt->v * lazy) % MOD;
        lt->marked = true;

        rt->lazy = (rt->lazy * lazy) % MOD;
        rt->v = (rt->v * lazy) % MOD;
        rt->marked = true;

        marked = false;
        lazy = 1ll;
    }

    void build(const vll& a) {
        if(l == r) {
            v = a[l];
            return;
        }

        ll m = (l + r) >> 1ll;

        lt = new Tree(l, m);
        rt = new Tree(m + 1ll, r);

        lt->build(a);
        rt->build(a);
    }

    void upd(ll ql, ll qr, ll qv) {
        if(ql > r || qr < l) return;

        push();

        if(ql == l && qr == r) {
            v = (v * qv) % MOD;
            lazy = (lazy * qv) % MOD;
            marked = true;

            return;
        }

        ll m = (l + r) >> 1ll;

        lt->upd(ql, min(qr, m), qv);
        rt->upd(max(ql, m + 1ll), qr, qv);
    }

    ll qry(ll i) {
        if(l == r) return v;
        return (i <= ((l + r) >> 1ll) ? lt : rt)->qry(i);
    }
};

inline vll invert_perm(const vll& a) {
    ll n = a.size();

    vll res(n, 0ll);

    for(ll i = 0ll; i < n; i++) {
        res[a[i]] = i;
    }

    return res;
}

void compute_subtree_sizes(ll i, ll pr, vll& sz, const vll& dfs_order, const vvll& adj) {
    sz[dfs_order[i]] = 1ll;

    for(ll j : adj[i]) {
        if(j == pr) continue;

        compute_subtree_sizes(j, i, sz, dfs_order, adj);
        sz[dfs_order[i]] += sz[dfs_order[j]];
    }
}

bool is_in_subtree(ll i, ll pr, const vll& dfs_order, const vll& sz) {
    return sz[dfs_order[pr]] >= dfs_order[i] - dfs_order[pr] + 1ll && dfs_order[i] >= dfs_order[pr];
}

void printv(const vll& a) {
    for(ll v : a) cerr << v << " ";
    cerr << endl;
}

/*
input: a target node (in original indexing) and a target depth
output: two node indices (in BFS order) 

first node invariant: l = set of all nodes such that dep[i] < dep_targ or dep[i] == dep_targ and dfs_order[i] < dfs_order[targ]
r is the first node

last node invariant: r = set of all nodes such that dep[i] > dep_targ or dep[i] == dep_targ and not within tree and sz[dfs_order[targ]] < dfs_order[i] - dfs_order[targ] + 1ll
l is the last node
*/
pll bounding_inds(ll targ, ll dep_targ, const vll& dep, const vll& dfs_order, const vll& bfs_order, const vll& bfs_order_inv, const vll& sz) {
    targ = bfs_order[targ];
    ll dfs_order_targ = dfs_order[bfs_order_inv[targ]];

    ll l = -1ll, r = dfs_order.size();
    
    while(r - l > 1ll) {
        ll m = (l + r) >> 1ll;

        if(dep[bfs_order_inv[m]] < dep_targ || (dep[bfs_order_inv[m]] == dep_targ && dfs_order[bfs_order_inv[m]] < dfs_order_targ)) l = m;
        else r = m;
    }

    ll fst = r;

    l = -1ll, r = dfs_order.size();


    ll targ_sz = sz[dfs_order_targ];

    while(r - l > 1ll) {
        ll m = (l + r) >> 1ll;

        if(dep[bfs_order_inv[m]] > dep_targ || (dep[bfs_order_inv[m]] == dep_targ && targ_sz < dfs_order[bfs_order_inv[m]] - dfs_order_targ + 1ll)) r = m;
        else l = m;
    }

    return {fst, l};
}

ll dfs_timer = 0ll;

void comp_dfs_order(ll i, ll pr, vll& dfs_order, const vvll& adj) {
    dfs_order[i] = dfs_timer++;

    for(ll j : adj[i]) {
        if(j == pr) continue;

        comp_dfs_order(j, i, dfs_order, adj);
    }
}

int main() {
    ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    
    // inputting graph information
    ll n;
    cin >> n >> MOD;

    vvll adj(n, vll());

    for(ll i = 0ll; i < n - 1ll; i++) {
        ll a, b;
        cin >> a >> b;

        a--; b--;

        adj[a].pb(b);
        adj[b].pb(a);
    }

    // computing dfs order, parents, and depth
    vll depths(n, -1ll);
    vll dfs_order(n, 0ll);

    stack<pll> s;
    s.push({0ll, 0ll});

    vb vis(n, false);
    vll par(n, -1ll);

    while(!s.empty()) {
        auto [i, pr] = s.top();
        s.pop();

        if(vis[i]) continue;
        vis[i] = true;
        par[i] = pr;
        depths[i] = depths[pr] + 1ll;

        for(ll j : adj[i]) {
            s.push({j, i});
        }
    }

    comp_dfs_order(0ll, 0ll, dfs_order, adj);

    // computing bfs order
    vll bfs_order(n, 0ll);

    queue<ll> q;

    ll bfs_timer = 0ll;

    for(ll i = 0ll; i < n; i++) vis[i] = false;

    q.push(0ll);

    while(!q.empty()) {
        ll i = q.front();

        q.pop();

        if(vis[i]) continue;
        vis[i] = true;
        bfs_order[i] = bfs_timer++;

        for(ll j : adj[i]) {
            q.push(j);
        }
    }

    vll bfs_order_inv = invert_perm(bfs_order);

    // computing initial height array
    vll h(n, 0ll);

    vll h_bfs_ind(n, 0ll);

    for(ll& v : h) cin >> v;

    for(ll i = 0ll; i < n; i++) {
        h_bfs_ind[bfs_order[i]] = h[i];
    }

    // building the segment tree
    Tree tr(0ll, n - 1ll);
    tr.build(h_bfs_ind);

    // computing subtree sizes
    vll sz(n, 0ll);
    compute_subtree_sizes(0ll, 0ll, sz, dfs_order, adj);


    // testing stuff
    // printv(dfs_order);
    // printv(sz);
    // printv(depths);
    // printv(bfs_order);
    // printv(bfs_order_inv);
    // pll res = bounding_inds(13ll, 2ll, depths, dfs_order, bfs_order, bfs_order_inv, sz);

    // cerr << res.first << " " << res.second << endl;

    // answering queries
    ll num_q;
    cin >> num_q;

    while(num_q--) {
        int qtype;
        cin >> qtype;

        if(qtype == 1) {
            // update
            ll x, d, w;
            cin >> x >> d >> w;
            x--;

            ll cur_targ = x;
            ll cur_dep_targ = depths[cur_targ];
            ll cur_dep_upd = cur_dep_targ + d;

            while(cur_dep_upd >= cur_dep_targ && cur_dep_upd >= 0ll) {
                auto [fst1, snd1] = bounding_inds(cur_targ, cur_dep_upd, depths, dfs_order, bfs_order, bfs_order_inv, sz);

                if(fst1 <= snd1) {
                    tr.upd(fst1, snd1, w);
                }

                cur_dep_upd--;

                if(cur_dep_upd < cur_dep_targ || cur_dep_upd < 0ll) break;

                auto [fst2, snd2] = bounding_inds(cur_targ, cur_dep_upd, depths, dfs_order, bfs_order, bfs_order_inv, sz);

                if(fst2 <= snd2) {
                    tr.upd(fst2, snd2, w);
                }

                cur_dep_upd--;
                cur_dep_targ--;
                cur_targ = par[cur_targ];
            }
        } else {
            // query
            ll x;
            cin >> x;
            x--;

            cout << tr.qry(bfs_order[x]) << "\n";
        }
    }

    cout << flush;

    return 0;
}

• Subtask 1: 0 / 3
• Subtask 2: 0 / 9
• Subtask 3: 0 / 29
• Subtask 4: 0 / 12
• Subtask 5: 0 / 30
• Subtask 6: 0 / 17