Submission #544837

#	Time	Username	Problem	Language	Result	Execution time	Memory
544837		MarcoMeijer	Sprinkler (JOI22_sprinkler)	C++14	100 / 100	2914 ms	166376 KiB

sprinkler

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int N = 2e5+10;

int n;
ll mod;
vector<int> adj[N];
int h[N];

// segment tree
void modify(vector<int>& seg, int l, int r, ll v) {
  int m = seg.size()/2;
  r = min(r, m-1);
  for (l += m, r += m+1; l < r; l /= 2, r/=2) {
    if (l&1) {
      seg[l] = ((ll)seg[l] * v)%mod;
      l++;
    }
    if (r&1) {
      r--;
      seg[r] = ((ll)seg[r] * v)%mod;
    }
  }
}
int query(vector<int>& seg, int p) {
  int res = 1;
  int m = seg.size()/2;
  for (p += m; p > 0; p /= 2) res = ((ll)res * (ll)seg[p])%mod;
  return res;
}

// centroid decompostion
int sz[N], rem[N];
vector<tuple<int,int,int>> centroids[N];
vector<vector<int>> seg[N];
int dfsSz(int u, int p=-1) {
  sz[u] = 1;
  for (int v : adj[u]) {
    if (rem[v] || v == p) continue;
    sz[u] += dfsSz(v, u);
  }
  return sz[u];
}
int dfsCent(int u, int p, int totSz) {
  for (int v : adj[u])
    if (!rem[v] && v != p)
      if (sz[v]*2 > totSz)
        return dfsCent(v, u, totSz);
  return u;
}
int dfs(int u, int p, int cent, int sub, int d) {
  centroids[u].push_back({cent, sub, d});
  int depth = 0;
  for (int v : adj[u])
    if (!rem[v] && v != p)
      depth = max(depth, dfs(v, u, cent, sub, d + 1));
  return depth + 1;
}
void handleCentroid(int u) {
  int sub = 0;
  int depth = 0;
  for (int v : adj[u]) {
    if (rem[v]) continue;
    int subDepth = dfs(v, u, u, sub, 1);
    depth = max(depth, subDepth+1);
    sub++;
  }
  depth = min(depth, 40);
  seg[u].resize(depth+1);
  centroids[u].push_back({u, -1, 0});
  for (int i=0; i<=depth; i++)
    seg[u][i].assign((sub+2)*2, 1);
}
void centroidDecomposition(int u) {
  dfsSz(u);
  int cent = dfsCent(u, -1, sz[u]);
  handleCentroid(cent);
  rem[cent] = 1;
  for (int v : adj[cent])
    if (!rem[v])
      centroidDecomposition(v);
}

// operations
const int SQ = 7;
void updateVertex(int u, int r, int v) {
  for (auto p : centroids[u]) {
    int cent, sub, d;
    tie(cent, sub, d) = p;
    int mx = min<int>(r-d+1, seg[cent].size());
    for (int i=0; i<mx; i++) {
      while (i + SQ < mx) i += SQ;
      if (sub == -1) {
        modify(seg[cent][i], 0, 1e9, v);
      } else {
        modify(seg[cent][i], 0, sub-1, v);
        modify(seg[cent][i], sub+1, 1e9, v);
      }
    }
  }
}
int queryVertex(int u) {
  int res = 1;
  auto add = [&](int x) {
    res = ((ll)res * (ll)x)%mod;
  };

  for (auto p : centroids[u]) {
    int cent, sub, d;
    tie(cent, sub, d) = p;
    if (d < seg[cent].size())
      add(query(seg[cent][d], sub));
    for (int i=0; i<seg[cent].size(); i+=SQ)
      if (i > d)
        add(query(seg[cent][i], sub));
  }
  return res;
}

int main() {
  // read graph
  cin >> n >> mod;
  for (int i=0; i<n-1; i++) {
    int u, v; cin >> u >> v;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }

  // centroid decomposition
  centroidDecomposition(1);

  // initial state
  for (int i=1; i<=n; i++) cin >> h[i];
  for (int i=1; i<=n; i++) updateVertex(i, 0, h[i]);

  // answer queries
  int q; cin >> q;
  while (q--) {
    int t; cin >> t;
    if (t == 1) {
      int x, d, w; cin >> x >> d >> w;
      updateVertex(x, d, w);
    }
    if (t == 2) {
      int x; cin >> x;
      cout << queryVertex(x) << endl;
    }
  }
}

Compilation message (stderr)

sprinkler.cpp: In function 'int queryVertex(int)':
sprinkler.cpp:114:11: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  114 |     if (d < seg[cent].size())
      |         ~~^~~~~~~~~~~~~~~~~~
sprinkler.cpp:116:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  116 |     for (int i=0; i<seg[cent].size(); i+=SQ)
      |                   ~^~~~~~~~~~~~~~~~~

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