Submission #1241069

#	Time	Username	Problem	Language	Result	Execution time	Memory
1241069		ProtonDecay314	Sprinkler (JOI22_sprinkler)	C++20	41 / 100	4094 ms	75176 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) {
//         push();
//         if(l == r) return v;
//         return (i <= ((l + r) >> 1ll) ? lt : rt)->qry(i);
//     }
// };

inline ll lt(ll i) {
    return i << 1ll;
}

inline ll rt(ll i) {
    return (i << 1ll) | 1ll;
}

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

inline void tree_construct(ll n, ll i, ll a_l, ll a_r, vll& l, vll& r, vll& v, vll& lazy, vb& marked, vb& isempty) {
    l[i] = a_l;
    r[i] = a_r;
    lazy[i] = 1ll;
    marked[i] = false;
    isempty[i] = false;
}

inline void push(ll n, ll i, vll& l, vll& r, vll& v, vll& lazy, vb& marked, vb& isempty) {
    if(!marked[i]) return;

    if(l[i] == r[i]) {
        marked[i] = false;
        lazy[i] = 1ll;
        return;
    }

    lazy[lt(i)] = (lazy[lt(i)] * lazy[i]) % MOD;
    v[lt(i)] = (v[lt(i)] * lazy[i]) % MOD;
    marked[lt(i)] = true;

    lazy[rt(i)] = (lazy[rt(i)] * lazy[i]) % MOD;
    v[rt(i)] = (v[rt(i)] * lazy[i]) % MOD;
    marked[rt(i)] = true;

    marked[i] = false;
    lazy[i] = 1ll;
}

inline void build(ll n, ll i, const vll& a, vll& l, vll& r, vll& v, vll& lazy, vb& marked, vb& isempty) {
    if(l[i] == r[i]) {
        v[i] = a[l[i]];
        return;
    }

    ll m = (l[i] + r[i]) >> 1ll;

    tree_construct(n, lt(i), l[i], m, l, r, v, lazy, marked, isempty);
    tree_construct(n, rt(i), m + 1ll, r[i], l, r, v, lazy, marked, isempty);

    build(n, lt(i), a, l, r, v, lazy, marked, isempty);
    build(n, rt(i), a, l, r, v, lazy, marked, isempty);
}

inline void upd(ll n, ll i, ll ql, ll qr, ll qv, vll& l, vll& r, vll& v, vll& lazy, vb& marked, vb& isempty) {
    if(ql > r[i] || qr < l[i]) return;

    push(n, i, l, r, v, lazy, marked, isempty);

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

        return;
    }

    ll m = (l[i] + r[i]) >> 1ll;

    upd(n, lt(i), ql, min(qr, m), qv, l, r, v, lazy, marked, isempty);
    upd(n, rt(i), max(ql, m + 1ll), qr, qv, l, r, v, lazy, marked, isempty);
}

inline ll qry(ll n, ll i, ll qi, vll& l, vll& r, vll& v, vll& lazy, vb& marked, vb& isempty) {
        push(n, i, l, r, v, lazy, marked, isempty);

        if(l[i] == r[i]) return v[i];
        return qry(n, (qi <= ((l[i] + r[i]) >> 1ll) ? lt(i) : rt(i)), qi, l, r, v, lazy, marked, isempty);
}

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
    vll l(n << 2ll, 0ll);
    vll r(n << 2ll, 0ll);
    vll v(n << 2ll, 1ll);
    vll lazy(n << 2ll, 1ll);
    vb marked(n << 2ll, false);
    vb isempty(n << 2ll, true);

    tree_construct(n, 1ll, 0ll, n - 1ll, l, r, v, lazy, marked, isempty);
    build(n, 1ll, h_bfs_ind, l, r, v, lazy, marked, isempty);

    // 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) {
                    upd(n, 1ll, fst1, snd1, w, l, r, v, lazy, marked, isempty);
                }

                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) {
                    upd(n, 1ll, fst2, snd2, w, l, r, v, lazy, marked, isempty);
                }

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

            cout << qry(n, 1ll, bfs_order[x], l, r, v, lazy, marked, isempty) << "\n";
        }
    }

    cout << flush;

    return 0;
}

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