E. Mahmoud and Ehab and the xor-MST | Codeforces Round #473 (Div. 2)

まず最上位ビットについて考えてみよう。最上位ビットが異なる頂点同士を結ぶと、ほかの全てのビットが異なる頂点同士を結ぶよりも高コストなので、できるだけ少なく済ませたい。そこで0b0???のグループと0b1???のグループに分けてそれぞれ木になっているとしよう。このとき、0b00...0と0b10...0を辺で結べばコストは最上位ビット1回分だけで済む。同様の問題を別れた部分木についても解く。つまり上位d-1桁が等しい値についてd桁目が異なる場合、d桁目で1本だけ辺を結べば済む。あとは、上位d-1桁が等しい値であるが、これは桁DPで数えていく。

dp[i][j] := i桁目まで見てj(0: n-1と一致, 1: n-1より小さい)ような上位i桁の数はいくつあるか
ll dp[51][2];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout << fixed << setprecision(20);

    ll n;
    cin >> n;
    n--;

    ll ans = 0;
    dp[0][0] = 1;
    rep(i, 50) {
        int k = 50 - 1 - i;
        int d = n >> k & 1;
        ll cost = 1ll << k;
        rep(les, 2) {
            ll x = dp[i][les];
            if (!x)continue;
            if (les) {
                ans += x * cost;
                dp[i + 1][1] += x * 2;
            } else {
                if (d == 1) {
                    ans += x * cost;
                    dp[i + 1][1] += x;
                }
                dp[i + 1][0] += x;
            }
        }
    }
    cout << ans << endl;
}