E2. PermuTree (hard version) TLE

Debugging Help
#1

Problem Info

E2. PermuTree (hard version)

Question

TLE on n ≥ 1E5

What I’ve Tried

My Work

Code

#include <bits/stdc++.h>

using namespace std;

const int N = 1E6;

long long ans = 0;

vector<int> adj[N];
int sz[N]; // subtreeSz
// bitset knapsack DP
template <int len = 1> void subsetSum(int n, map<int, int> &szCnt) {
  if (n >= len) { // pseudo var len bitset
    subsetSum<min(2 * len, N)>(n, szCnt);
    return;
  }

  bitset<len> dp; // dp[sz]: can sum(some kids subtreeSz) == sz
  dp[0] = 1;
  // {s, 12} -> {{s,1}, {s,2}, {s,3}, {s,4}, {s,5}}
  for (auto &[s, cnt] : szCnt) {
    for (int i = 1; i <= cnt; cnt -= i, i *= 2) {
      dp |= (dp << (i * s));
    }
    // till here {s, 1}, {s, 2}, {s, 3}, { s, 4 }
    if (cnt == 0) {
      continue;
    }

    dp |= (dp << (cnt * s)); // {s, 5}
  }

  long long res = 0;

  for (int i = 1; i < n; i++) {
    if (!dp[i]) {
      continue;
    }

    res = max(res, i * (n - i - 1LL)); // i kids val < me & rest kids val > me
  }

  ans += res;
}

void dfs(int me) {
  sz[me] = 1;
  int mxSz = 0; // biggest subtreeSz of kids

  map<int, int> szCnt; // kids' {subtreeSz, count}

  for (int kid : adj[me]) {
    dfs(kid);
    sz[me] += sz[kid];
    mxSz = max(mxSz, sz[kid]);
    szCnt[sz[kid]]++;
  }

  if (mxSz * 2 >= sz[me]) { // small to large merge trick for efficiency
    ans += (mxSz * (sz[me] - mxSz - 1LL));
    return;
  }

  subsetSum(sz[me], szCnt);
}

int main() {
  int n;
  scanf("%d", &n);

  for (int i = 1; i < n; i++) {
    int p;
    scanf("%d", &p);

    p--;

    adj[p].push_back(i); // par -> kid
  }

  dfs(0); // O(n√n)

  printf("%lld\n", ans);
}

Generator

#include <bits/stdc++.h>

using namespace std;

int rand(int a, int b) { return a + rand()%(b-a); }

int main() {
  int n;
  cin >> n;
  cout << n << endl;
  for (int i=2; i <= n; i++) { cout << rand(1,i) << ' '; }
  
  cout << endl;
 }

Logic

Editorial.