Submission #1245362

#	Time	Username	Problem	Language	Result	Execution time	Memory
1245362		ProtonDecay314	Sprinkler (JOI22_sprinkler)	C++20	100 / 100	882 ms	94372 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!

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
*/

void find_par(ll i, ll pr, vll& par, const vvll& adj) {
    par[i] = pr;
    
    for(ll j : adj[i]) {
        if(j == pr) continue;

        find_par(j, i, par, 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);
    }

    // find parents

    vll par(n, -1ll);

    find_par(0ll, 0ll, par, adj);

    // store updates
    const ll MAXD = 40ll;
    vvll upds(n, vll(MAXD + 1ll, 1ll));

    vll init(n, 0ll);
    for(ll i = 0ll; i < n; i++) cin >> init[i]; 

    // perform the updates
    ll q;
    cin >> q;

    while(q--) {
        int qt;
        cin >> qt;

        if(qt == 1) {
            // upd
            ll x, d, w;
            cin >> x >> d >> w;

            x--;
            ll cur_d = d;
            ll cur_n = x;

            while(cur_d >= 0ll) {
                if(par[cur_n] != cur_n) {
                    // "normal" update

                    upds[cur_n][cur_d] *= w;
                    upds[cur_n][cur_d] %= MOD;

                    if(cur_d > 0ll) {
                        upds[cur_n][cur_d - 1ll] *= w;
                        upds[cur_n][cur_d - 1ll] %= MOD;
                    }
                    cur_d--;
                } else {
                    // update to completion

                    for(ll i = 0ll; i <= cur_d; i++) {
                        upds[cur_n][i] *= w;
                        upds[cur_n][i] %= MOD;
                    }

                    break;
                }

                cur_n = par[cur_n];
            }
        } else {
            // qry
            ll x;
            cin >> x;
            x--;
            ll res = init[x];

            ll cur_n = x;
            ll cur_d = 0ll;

            while(cur_d <= MAXD) {
                res = (res * upds[cur_n][cur_d]) % MOD;

                cur_d++;

                if(par[cur_n] == cur_n) break;
                cur_n = par[cur_n];
            }

            cout << res << "\n";
        }
    }

    cout << flush;

    return 0;
}

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