QOJ.ac

QOJ

IDProblemSubmitterResultTimeMemoryLanguageFile sizeSubmit timeJudge time
#296034#4845. DFSHoMaMaOvO (Riku Kawasaki, Masaki Nishimoto, Yui Hosaka)#TL 108ms14560kbC++147.0kb2024-01-01 23:23:232024-01-01 23:23:23

Judging History

This is the latest submission verdict.

  • [2024-01-01 23:23:23]
  • Judged
  • Verdict: TL
  • Time: 108ms
  • Memory: 14560kb
  • [2024-01-01 23:23:23]
  • Submitted

answer

#include <cassert>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <bitset>
#include <complex>
#include <deque>
#include <functional>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

using Int = long long;

template <class T1, class T2> ostream &operator<<(ostream &os, const pair<T1, T2> &a) { return os << "(" << a.first << ", " << a.second << ")"; };
template <class T> ostream &operator<<(ostream &os, const vector<T> &as) { const int sz = as.size(); os << "["; for (int i = 0; i < sz; ++i) { if (i >= 256) { os << ", ..."; break; } if (i > 0) { os << ", "; } os << as[i]; } return os << "]"; }
template <class T> void pv(T a, T b) { for (T i = a; i != b; ++i) cerr << *i << " "; cerr << endl; }
template <class T> bool chmin(T &t, const T &f) { if (t > f) { t = f; return true; } return false; }
template <class T> bool chmax(T &t, const T &f) { if (t < f) { t = f; return true; } return false; }
#define COLOR(s) ("\x1b[" s "m")

////////////////////////////////////////////////////////////////////////////////
template <unsigned M_> struct ModInt {
  static constexpr unsigned M = M_;
  unsigned x;
  constexpr ModInt() : x(0U) {}
  constexpr ModInt(unsigned x_) : x(x_ % M) {}
  constexpr ModInt(unsigned long long x_) : x(x_ % M) {}
  constexpr ModInt(int x_) : x(((x_ %= static_cast<int>(M)) < 0) ? (x_ + static_cast<int>(M)) : x_) {}
  constexpr ModInt(long long x_) : x(((x_ %= static_cast<long long>(M)) < 0) ? (x_ + static_cast<long long>(M)) : x_) {}
  ModInt &operator+=(const ModInt &a) { x = ((x += a.x) >= M) ? (x - M) : x; return *this; }
  ModInt &operator-=(const ModInt &a) { x = ((x -= a.x) >= M) ? (x + M) : x; return *this; }
  ModInt &operator*=(const ModInt &a) { x = (static_cast<unsigned long long>(x) * a.x) % M; return *this; }
  ModInt &operator/=(const ModInt &a) { return (*this *= a.inv()); }
  ModInt pow(long long e) const {
    if (e < 0) return inv().pow(-e);
    ModInt a = *this, b = 1U; for (; e; e >>= 1) { if (e & 1) b *= a; a *= a; } return b;
  }
  ModInt inv() const {
    unsigned a = M, b = x; int y = 0, z = 1;
    for (; b; ) { const unsigned q = a / b; const unsigned c = a - q * b; a = b; b = c; const int w = y - static_cast<int>(q) * z; y = z; z = w; }
    assert(a == 1U); return ModInt(y);
  }
  ModInt operator+() const { return *this; }
  ModInt operator-() const { ModInt a; a.x = x ? (M - x) : 0U; return a; }
  ModInt operator+(const ModInt &a) const { return (ModInt(*this) += a); }
  ModInt operator-(const ModInt &a) const { return (ModInt(*this) -= a); }
  ModInt operator*(const ModInt &a) const { return (ModInt(*this) *= a); }
  ModInt operator/(const ModInt &a) const { return (ModInt(*this) /= a); }
  template <class T> friend ModInt operator+(T a, const ModInt &b) { return (ModInt(a) += b); }
  template <class T> friend ModInt operator-(T a, const ModInt &b) { return (ModInt(a) -= b); }
  template <class T> friend ModInt operator*(T a, const ModInt &b) { return (ModInt(a) *= b); }
  template <class T> friend ModInt operator/(T a, const ModInt &b) { return (ModInt(a) /= b); }
  explicit operator bool() const { return x; }
  bool operator==(const ModInt &a) const { return (x == a.x); }
  bool operator!=(const ModInt &a) const { return (x != a.x); }
  friend std::ostream &operator<<(std::ostream &os, const ModInt &a) { return os << a.x; }
};
////////////////////////////////////////////////////////////////////////////////

constexpr unsigned MO = 998244353;
using Mint = ModInt<MO>;

constexpr int LIM_INV = 800'010;
Mint inv[LIM_INV], fac[LIM_INV], invFac[LIM_INV];

void prepare() {
  inv[1] = 1;
  for (int i = 2; i < LIM_INV; ++i) {
    inv[i] = -((Mint::M / i) * inv[Mint::M % i]);
  }
  fac[0] = invFac[0] = 1;
  for (int i = 1; i < LIM_INV; ++i) {
    fac[i] = fac[i - 1] * i;
    invFac[i] = invFac[i - 1] * inv[i];
  }
}
Mint binom(Int n, Int k) {
  if (n < 0) {
    if (k >= 0) {
      return ((k & 1) ? -1 : +1) * binom(-n + k - 1, k);
    } else if (n - k >= 0) {
      return (((n - k) & 1) ? -1 : +1) * binom(-k - 1, n - k);
    } else {
      return 0;
    }
  } else {
    if (0 <= k && k <= n) {
      assert(n < LIM_INV);
      return fac[n] * invFac[k] * invFac[n - k];
    } else {
      return 0;
    }
  }
}


int N, R;
vector<int> C;
vector<int> A, B;

vector<pair<Int, int>> cus;
vector<int> ids;
vector<vector<int>> graph;


namespace brute {
vector<int> mns;
void dfs(int u, int p) {
  mns[u] = ids[u];
  for (const int v : graph[u]) if (p != v) {
    dfs(v, u);
    chmin(mns[u], mns[v]);
  }
}
Mint ans;
map<int, Mint> solve(int u, int p) {
  vector<pair<int, int>> xvs;
  for (const int v : graph[u]) if (p != v) {
    xvs.emplace_back(mns[v], v);
  }
  sort(xvs.begin(), xvs.end());
  const int len = xvs.size();
  map<int, Mint> tmp;
  for (int j = 0; j < len; ++j) {
    const int v = xvs[j].second;
    const auto res = solve(v, u);
    for (const auto &kv : res) {
      assert(xvs[j].first <= kv.first);
      int k;
      for (k = 0; k < len && xvs[k].first <= kv.first; ++k) if (j != k) {
        tmp[xvs[k].first] += ((k < j) ? (inv[k + 1] * inv[k + 2]) : (inv[k] * inv[k + 1])) * kv.second;
      }
      tmp[kv.first] += inv[k] * kv.second;
    }
  }
  map<int, Mint> ret;
  ret[ids[u]] += 1;
  for (const auto &kv : tmp) {
    ret[min(ids[u], kv.first)] += kv.second;
  }
  for (const auto &kv : ret) {
    ans += kv.second * cus[kv.first].first;
  }
/*
cerr<<"u = "<<u<<": xvs = "<<xvs<<endl;
cerr<<"u = "<<u<<": tmp = ";pv(tmp.begin(),tmp.end());
cerr<<"u = "<<u<<": ret = ";pv(ret.begin(),ret.end());
*/
  return ret;
}
Mint run() {
  mns.assign(N, -1);
  dfs(R, -1);
  ans = 0;
  solve(R, -1);
  return ans;
}
}  // brute



int main() {
  prepare();
  
  for (int numCases; ~scanf("%d", &numCases); ) { for (int caseId = 1; caseId <= numCases; ++caseId) {
    scanf("%d%d", &N, &R);
    --R;
    C.resize(N);
    for (int u = 0; u < N; ++u) {
      scanf("%d", &C[u]);
    }
    A.resize(N - 1);
    B.resize(N - 1);
    for (int i = 0; i < N - 1; ++i) {
      scanf("%d%d", &A[i], &B[i]);
      --A[i];
      --B[i];
    }
    
    cus.resize(N);
    for (int u = 0; u < N; ++u) {
      cus[u] = make_pair(C[u], u);
    }
    sort(cus.begin(), cus.end());
    ids.assign(N, -1);
    for (int x = 0; x < N; ++x) {
      ids[cus[x].second] = x;
    }
// cerr<<"cus = "<<cus<<endl;
    
    graph.assign(N, {});
    for (int i = 0; i < N - 1; ++i) {
      graph[A[i]].push_back(B[i]);
      graph[B[i]].push_back(A[i]);
    }
    
    const Mint ans = brute::run();
    printf("%u\n", ans.x);
    
#ifdef LOCAL
const Mint brt=brute::run();
cerr<<"brt = "<<brt<<endl;
#endif
  }
#ifndef LOCAL
  break;
#endif
  }
  return 0;
}

Details

Tip: Click on the bar to expand more detailed information

Test #1:

score: 100
Accepted
time: 8ms
memory: 13172kb

input:

4
1 1
1
3 3
3 3 4
3 1
3 2
6 1
5 2 4 1 3 6
1 2
1 6
2 3
2 4
4 5
5 1
5 4 3 2 1
1 2
1 3
3 4
3 5

output:

1
16
34
499122202

result:

ok 4 number(s): "1 16 34 499122202"

Test #2:

score: 0
Accepted
time: 92ms
memory: 14256kb

input:

7
5000 933
23306350 162661794 68618194 666430282 995855733 929210414 295740530 464135554 304211641 725090719 226242817 592655639 936895997 479520010 108891341 598601399 678169271 118406229 394867734 640888099 481066130 606481085 709600400 554804145 179044332 41718098 549318629 400214219 159098456 67...

output:

647896606
670593316
448857064
140431373
205960849
578484974
77271255

result:

ok 7 numbers

Test #3:

score: 0
Accepted
time: 108ms
memory: 14560kb

input:

7
5000 3103
370388267 486433577 320921400 202742370 718520472 895122554 359601184 337298427 146232648 830940586 826047977 229951976 919497287 67059290 843962706 777524684 196246825 682863698 122309598 743014832 349595264 964812795 660215381 963691135 376209511 759526115 822829377 916639371 802105594...

output:

371972695
888631092
915285114
181894451
144642157
385851099
551331307

result:

ok 7 numbers

Test #4:

score: 0
Accepted
time: 95ms
memory: 14528kb

input:

7
5000 4968
717470184 105172657 308383390 593830267 706026427 861034695 423461838 60718197 988253654 231757749 690694353 162215609 461907090 804341675 874001367 811223777 714324378 97578063 704527271 140108861 363348590 28177209 760573466 522321229 868341986 917525621 541050526 287840331 295369628 8...

output:

249142844
829572814
771685997
331594252
670718819
285030671
941548421

result:

ok 7 numbers

Test #5:

score: -100
Time Limit Exceeded

input:

6
100000 64523
754457004 816672702 108228103 727565046 575355672 557029941 374900316 667662208 247774859 241366803 327704349 984466565 870042639 983615018 950320614 473363640 338609139 877917484 271995097 819820270 213958988 431503923 943718683 511580322 791201903 944289849 290330452 148342386 88563...

output:


result: