OI-codes/S2OJ/2021/data/chk.cpp

#include "testlib.h"

typedef long long int ll;

const int maxn = 35;

int pre[maxn], a[maxn];
ll scr[maxn];

struct Node {
  int id;
  Node *ls, *rs;

  Node(int x) : id(x) {
    ls = rs = nullptr;
  }
};
Node *node[maxn];

void build(int L, int R, int l, int r);
void dfs(const int u);

int main(int argc, char* argv[]){
  registerTestlibCmd(argc, argv);
  ll pans = ouf.readLong(), jans = ans.readLong();
  if (jans != pans) {
    quitf(_wa, "OvO, the ans in first line is %lld, but yours is %lld", jans, pans);
  }
  int n = inf.readInt();
  for (int i = 1; i <= n; ++i) {
    a[i] = inf.readInt();
    pre[i] = ouf.readInt();
  }
  build(1, n, 1, n);
  dfs(pre[1]);
  if (scr[pre[1]] == jans) {
    quitf(_ok, "Orz, your ans is perfectly correct!");
  } else {
    quitp(0.5, "OvO, your tree seems to be wrong QAQ");
  }
}

void dfs(const int u) {
  ll lv = 1, rv = 1;
  if (node[u]->ls) {
    dfs(node[u]->ls->id);
    lv = scr[node[u]->ls->id];
  }
  if (node[u]->rs) {
    dfs(node[u]->rs->id);
    rv = scr[node[u]->rs->id];
  }
  if ((node[u]->ls == nullptr) && (node[u]->rs == nullptr)) lv = 0;
  scr[u] = lv * rv + a[u];
}

void build(int L, int R, int l, int r) {
  if (L > R) return;
  node[pre[L]] = new Node(pre[L]);
  if (l == r) return;
  int rt = pre[L];
  int lsz = rt - l;
  if (lsz) {
    build(L + 1, L + lsz, l, rt - 1);
    node[rt]->ls = node[pre[L + 1]];
  }
  int rsz = r - rt;
  if (rsz) {
    build(R - rsz + 1, R, rt + 1, r);
    node[rt]->rs = node[pre[R - rsz + 1]];
  }
}