# E2. PermuTree (hard version) TLE

## Problem Info

E2. PermuTree (hard version)

TLE on n ≥ 1E5

## My Work

### Code

``````#include <bits/stdc++.h>

using namespace std;

const int N = 1E6;

long long ans = 0;

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--;

}

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;
}
``````