Tree Sum is Fun II (Editorial)

Problem Author:	Hasan Saleh (MathModel)
Prepared by:	Hasan Saleh (MathModel)
Edited by:	Kostiantyn Denysov

Let us first consider a simplified version of the problem where we are only allowed to move from a node to its parent.

This means that each node in the path must be the parent of the next one:

u_{i} = ⌊ \frac{u _{i + 1}}{2} ⌋ .

u_i = \left\lfloor \frac{u_{i+1}}{2} \right\rfloor.

In other words, we are allowed to construct a path by repeatedly moving from a node to its parent. The path must be strictly increasing, so we must construct it in reverse — starting from some node $v$ and repeatedly moving to its parent.

Let us define the ancestor path of a node $v$ as:

A (v) = {p^{0} (v), p^{1} (v), p^{2} (v), \dots},

A(v) = \{p^0(v), p^1(v), p^2(v), \dots\},

where $p^{k} (v)$ p^k(v) denotes the $k$ -th ancestor of $v$ , and the path terminates when $p^{k} (v) < p^{k + 1} (v)$ p^k(v) < p^{k+1}(v).

However, in this problem, we are allowed to end the path at any point, not necessarily at the root. So for each node $v$ , we can choose a prefix of its ancestor chain of any length $l \geq 1$ l \ge 1, and check whether the sum of that prefix equals $n$ .

Our goal becomes: for how many pairs $(v, l)$ (v, l) does the prefix of length $l$ of the ancestor chain of $v$ sum to $n$ ?

Let us now consider how to find valid paths when the length $l$ of the path is fixed.

As we are only moving from a node to its parent at each step, a path of length $l$ is uniquely determined by its starting node $v$ :

v, ⌊ \frac{v}{2} ⌋, ⌊ \frac{v}{4} ⌋, \dots, ⌊ \frac{v}{2 ^{l - 1}} ⌋ .

v, \left\lfloor \frac{v}{2} \right\rfloor, \left\lfloor \frac{v}{4} \right\rfloor, \dots, \left\lfloor \frac{v}{2^{l-1}} \right\rfloor.

Let us denote the sum of this path as $S (v, l)$ S(v, l). Then:

S (v, l) = i = 0 \sum l - 1 ⌊ \frac{v}{2 ^{i}} ⌋ .

S(v, l) = \sum_{i = 0}^{l - 1} \left\lfloor \frac{v}{2^i} \right\rfloor.

Our goal is, for each fixed $l$ , to find how many integers $v$ satisfy $S (v, l) = n$ S(v, l) = n.

It is easy to observe that $S (v, l)$ S(v, l) is a strictly increasing function of $v$ . Indeed, increasing $v$ by $1$ increases each term $⌊ \frac{v}{2 ^{i}} ⌋$ \left\lfloor \frac{v}{2^i} \right\rfloor by at least $0$ , and increases the first term (i.e., $v$ itself) by $1$ .

Therefore, for fixed $l$ , the function $S (v, l)$ S(v, l) is strictly increasing, which allows us to efficiently find value of $v$ such that $S (v, l) = n$ S(v, l) = n using binary search (if it exists).

Although binary search works well in theory, we can avoid it entirely by using a more analytical approach.

Let us fix the path length $l$ , and consider how the value of $S (v, l)$ S(v, l) depends on the binary representation of $v$ .

We observe that the sum

S (v, l) = i = 0 \sum l - 1 ⌊ \frac{v}{2 ^{i}} ⌋

S(v, l) = \sum_{i=0}^{l-1} \left\lfloor \frac{v}{2^i} \right\rfloor

depends on the lowest $l$ bits of $v$ , as well as the higher bits above position $l$ .

Let us write $v$ in binary as a concatenation of two parts:

v = 2^{l} \cdot x + y,

v = 2^l \cdot x + y,

where:

$y$ consists of the lowest $l$ bits of $v$ (i.e., $0 \leq y < 2^{l}$ 0 \le y < 2^l),
$x$ is the prefix — the higher bits of $v$ .

We now analyze $S (v, l)$ S(v, l) in terms of $x$ and $y$ . Since the operation $⌊ \frac{v}{2 ^{i}} ⌋$ \left\lfloor \frac{v}{2^i} \right\rfloor becomes:

⌊ \frac{2 ^{l} x + y}{2 ^{i}} ⌋ = 2^{l - i} x + ⌊ \frac{y}{2 ^{i}} ⌋,

\left\lfloor \frac{2^l x + y}{2^i} \right\rfloor = 2^{l - i} x + \left\lfloor \frac{y}{2^i} \right\rfloor,

we can rewrite the entire sum as:

S (v, l) = x \cdot i = 0 \sum l - 1 2^{l - i} + i = 0 \sum l - 1 ⌊ \frac{y}{2 ^{i}} ⌋ .

S(v, l) = x \cdot \sum_{i = 0}^{l - 1} 2^{l - i} + \sum_{i = 0}^{l - 1} \left\lfloor \frac{y}{2^i} \right\rfloor.

The first sum is a geometric series:

i = 0 \sum l - 1 2^{l - i} = 2^{l + 1} - 1,

\sum_{i = 0}^{l - 1} 2^{l - i} = 2^{l + 1} - 1,

The second sum can be simplified as well. It is known that:

i = 0 \sum l - 1 ⌊ \frac{y}{2 ^{i}} ⌋ = 2 y - popcount (y),

\sum_{i = 0}^{l - 1} \left\lfloor \frac{y}{2^i} \right\rfloor = 2y - \text{popcount}(y),

where $popcount (y)$ denotes the number of ones in the binary representation of $y$ .

Therefore, the final expression becomes:

S (v, l) = x \cdot (2^{l + 1} - 1) + (2 y - popcount (y)) .

S(v, l) = x \cdot (2^{l + 1} - 1) + (2y - \text{popcount}(y)).

If we brute force $l$ and $popcount (y)$ we can find $x$ and $y$ as $y$ is not bigger then $2^{l}$ 2^l. Therefore $2 y = (n + popcount (y)) mod (2^{l + 1} - 1)$ 2y=(n+\text{popcount}(y)) \mod (2^{l + 1} - 1).

Handling Uncle Transitions

Once we find all valid starting nodes $v$ for a path of length $l$ such that $S (v, l) = n$ S(v, l) = n using only parent transitions, we need to account for the additional possibility of using uncle transitions in the path.

Recall that at each step in the path, we are allowed to move from a node $u$ to either its parent $p (u)$ or its uncle $q (u)$ . Importantly, both $p (u)$ and $q (u)$ share the same grandparent, so their values differ by $1$ :

∣ p (u) - q (u) ∣ = 1.

|p(u) - q(u)| = 1.

Therefore, in any path of length $l$ , each node (except the last one) can be either a parent or an uncle of the next node. If we fix the sequence of parent-child relationships (i.e., a path built only from parent transitions), we can consider replacing some of the steps with uncle transitions.

Each time we replace a parent with an uncle, we are changing the corresponding value in the path by $\pm 1$ . The effect on the total sum of the path is:

$+ 1$ if we choose the uncle with larger value than the parent,
$- 1$ if we choose the smaller one.

Which one is chosen depends on the parity of the parent — specifically, the second least significant bit in the binary representation of the node.

Let us denote:

For a node $u$ , its parent is $p = ⌊ u /2 ⌋$ p = \lfloor u/2 \rfloor.
Its uncle is $q = p \oplus 1$ q = p \oplus 1, since they are siblings.

Hence, replacing a parent with an uncle shifts the value by $\pm 1$ , and the direction of change depends on whether $p$ is even or odd, i.e., on the second lowest bit of $u$ .

Thus, if we take a path of length $l$ consisting only of parent transitions and total sum $n^{'}$ n', we can generate multiple distinct paths by flipping parent steps to uncle steps. Each flip adds or subtracts $1$ from the total sum depending on the binary structure of the path. This allows us to cover a range of sums centered at $n^{'}$ n'.

To find all values of $n$ for which such a path exists, we can:

Fix a base path with sum $n^{'}$ n'.
Simulate all $2^{l - 1}$ 2^{l-1} combinations of using either a parent or an uncle at each step.
For each such combination, compute the resulting total sum and count the number of ways it matches $n$ .

However, we can do this smarter — we will describe this in the next section.

Recall that for each valid base path of length $l$ using only parent transitions, we computed its total sum $n^{'}$ n'. Our goal is to determine, for each input $n$ , how many ways exist to obtain a path with sum exactly $n$ by selectively replacing parent steps with uncle steps.

Each replacement of a parent by its uncle changes the corresponding node in the path by either $+ 1$ or $- 1$ , depending on whether the uncle is larger or smaller than the parent.

The direction of this change depends on the parity of the parent — in particular, on the second least significant bit of the current node.
Specifically:
- If the second lowest bit is $0$ (i.e., parent is even), switching to uncle adds $+ 1$ to the sum.
- If the second lowest bit is $1$ (i.e., parent is odd), switching to uncle subtracts $1$ from the sum.

Let $c_{+}$ c_+ be the number of steps where switching to an uncle increases the sum, and $c_{-}$ c_- be the number of steps where it decreases the sum. Since the path has $l$ transitions, we have:

c_{+} + c_{-} \leq l .

c_+ + c_- \le l.

If we flip $x$ increasing steps and $y$ decreasing steps, the new sum becomes:

n = n^{'} + x - y .

n = n' + x - y.

Hence, the difference between $n$ and $n^{'}$ n' must be in the range $[- l, l]$ [-l, l], and for each such possible difference, we can try all $(x, y)$ pairs satisfying:

x - y = n - n^{'}, 0 \leq x \leq c_{+}, 0 \leq y \leq c_{-} .

x - y = n - n', \quad 0 \le x \le c_+, \quad 0 \le y \le c_-.

Now, for fixed $c_{+}$ c_+ and $c_{-}$ c_- (which we can count by examining the bits of the path), the number of ways to flip $x$ out of $c_{+}$ c_+ and $y$ out of $c_{-}$ c_- is:

(x c _{+}) \cdot (y c _{-}) .

\binom{c_+}{x} \cdot \binom{c_-}{y}.

Therefore, for each valid base path (determined by $v$ and $l$ ), and for each possible offset $d = n - n^{'} \in [- l, l]$ d = n - n' \in [-l, l], we compute the number of flips $(x, y)$ such that $x - y = d$ x - y = d, and sum over all such combinations.

This yields the total number of paths of length $l$ with sum exactly $n$ , including those with uncle transitions.

For a given base path, we want to compute the number of ways to flip $x$ transitions from the $c_{+}$ c_+ group and $y$ transitions from the $c_{-}$ c_- group such that:

x - y = d = n - n^{'},

x - y = d = n - n',

and $0 \leq x \leq c_{+}$ 0 \le x \le c_+, $0 \leq y \leq c_{-}$ 0 \le y \le c_-.

Instead of iterating over all $(x, y)$ pairs explicitly, we can precompute, for each possible triple $(d, c_{+}, c_{-})$ (d, c_+, c_-), the number of ways to choose $(x, y)$ such that $x - y = d$ x - y = d. This reduces the computation to simple table lookups and binomial coefficients.

Final Optimization and Complexity

To solve the problem efficiently, we iterate over all possible path lengths $l \leq 60$ l \le 60 and all differences $d = n - n^{'}$ d = n - n' in the range $[- l, l]$ [-l, l].

For each $(l, d)$ we compute the corresponding value $v$ using the analytical formula for $S (v, l)$ S(v, l) or binary search, and check whether $v$ is valid.

Then, we determine how many transitions in the path can increase or decrease the total sum when flipping a parent to an uncle, denoted as $c_{+}$ c_+ and $c_{-}$ c_- respectively.

Since we only need the difference $x - y = d$ x - y = d, not the exact values of $x$ and $y$ , we precompute the number of ways to achieve each $d$ using dynamic programming over $(c_{+}, c_{-}, d)$ (c_+, c_-, d).

This results in a total time complexity of $O (lo g^{3} n)$ O(\log^3 n), which is fast enough.

//#pragma GCC optimize("Ofast", "unroll-loops")
//#pragma GCC target("sse", "sse2", "sse3", "ssse3", "sse4")

#include <bits/stdc++.h>

#define all(a) a.begin(),a.end()
#define len(a) (int)(a.size())
#define mp make_pair
#define pb push_back
#define fi first
#define se second

using namespace std;

typedef pair<int, int> pii;
typedef long long ll;
typedef long double ld;
template<class T>
using vec = vector<T>;

template<typename T>
bool umin(T &a, T b) {
    if (b < a) {
        a = b;
        return true;
    }
    return false;
}
template<typename T>
bool umax(T &a, T b) {
    if (a < b) {
        a = b;
        return true;
    }
    return false;
}

#ifdef KoRoVa
#define DEBUG for (bool _FLAG = true; _FLAG; _FLAG = false)
#define LOG(...) print(#__VA_ARGS__" ::", __VA_ARGS__) << endl
template <class ...Ts> auto &print(Ts ...ts) { return ((cerr << ts << " "), ...); }
#else
#define DEBUG while (false)
#define LOG(...)
#endif

const int max_n = 100, inf = 1000111222;


const int mod = 1e9 + 7;


inline void inc(int &x, int y) {
    x += y;
    if (x >= mod) {
        x -= mod;
    }
}

inline void dec(int &x, int y) {
    x -= y;
    if (x < 0) {
        x += mod;
    }
}

inline int neg (int x) {
    return x ? mod - x : 0;
}

inline int mul(int x, int y) {
    return (1LL * x * y) % mod;
}

int power(int x, int y) {
    if (y < 0) {
        y += mod - 1;
    }
    if (y == 0) {
        return 1;
    }
    if (y % 2 == 0) {
        return power(mul(x, x), y / 2);
    }
    return mul(x, power(x, y - 1));
}

inline int inv(int x) {
    return power(x, mod - 2);
}


// new mint part


inline int divide (int x, int y) {
    return mul(x, inv(y));
}

struct mint {
#define oper_apply(op, f) mint& operator op (const mint &x) {f(val, x.val); return *this;}
#define oper_apply_const(op, f) mint operator op (const mint &x) const {mint tmp = *this; f(tmp.val, x.val); return tmp;}
#define oper_return(op, f) mint& operator op (const mint &x) {val = f(val, x.val); return *this;}
#define oper_return_const(op, f) mint operator op (const mint &x) const {return mint(f(val, x.val));}

    int val;

    mint(int x = 0) : val(x) {}

    oper_apply(+=, inc);
    oper_apply_const(+, inc);
    oper_apply(-=, dec);
    oper_apply_const(-, dec);
    oper_return(*=, mul);
    oper_return_const(*, mul);
    oper_return(/=, divide);
    oper_return_const(/, divide);
    oper_return(^=, power);
    oper_return_const(^, power);
};

const int max_f = max_n;
const int max_rf = max_f;
static_assert(max_f >= 1 && max_rf >= 1);

mint f[max_f], rf[max_rf];

bool frf_calculated = []() {
    f[0] = 1;
    for (int i = 1; i < max_f; ++i) {
        f[i] = f[i - 1] * mint(i);
    }
    mint x = f[min(max_f, max_rf) - 1];
    for (int i = max_f; i < max_rf; ++i) {
        x *= mint(i);
    }
    rf[max_rf - 1] = inv(x.val);
    for (int i = max_rf - 2; i >= 0; --i) {
        rf[i] = rf[i + 1] * mint(i + 1);
    }
    return true;
}();

inline mint A (int n, int k) {
    if (k < 0 || k > n) return 0;
    return f[n] * rf[n - k];
}

inline mint C (int n, int k) {
    if (k < 0 || k > n) return 0;
    return A(n, k) * rf[k];
}


mt19937 rng(228);
template<typename T = int>
inline T randll(T l = INT_MIN, T r = INT_MAX) {
    return uniform_int_distribution<T>(l, r)(rng);
}

inline ld randld(ld l = INT_MIN, ld r = INT_MAX) {
    return uniform_real_distribution<ld>(l, r)(rng);
}

const int debug = 0;

inline mint stupid (int n) {
    mint ans = 0;
    vector <map <int, mint> > dp(n + 1);
    for (int i = n; i > 0; i--) {
        dp[i][i] += 1;
        if (i > 3) {
            int go1 = i / 2;
            int go2 = go1 ^ 1;
            for (auto &j: dp[i]) {
                if (j.first + go1 <= n) {
                    dp[go1][j.first + go1] += j.second;
                }
                if (j.first + go2 <= n) {
                    dp[go2][j.first + go2] += j.second;
                }
            }
        }
        else if (i > 1) {
            int go = 1;
            for (auto &j: dp[i]) {
                if (j.first + go <= n) {
                    dp[go][j.first + go] += j.second;
                }
            }
        }
        if (dp[i].count(n)) {
            ans += dp[i][n];
        }
    }
    return ans;
}

mint dp[max_n * 2][max_n][max_n];

inline void test_case () {
    ll n;cin>>n;
    mint ans = 1;

    for (int l = 1; l <= 60; l++) {
        ll c = (1ll << (l + 1)) - 1;
        for (int diff = -l; diff <= l; diff++) {
            ll final = n + diff;
            if (final <= 0) {
                continue;
            }
            ll val = final % c;

            for (int cnt = 0; cnt <= 60; cnt++) {
                ll gg = val + cnt;
                if (gg & 1ll) {
                    continue;
                }
                gg /= 2;
                if (__builtin_popcountll(gg) != cnt) {
                    continue;
                }
                ll x = (final + cnt - 2 * gg);
                assert(x % c == 0);
                x /= c;
//                    LOG(l, gg, x, c);
                if (x) {
                    ll cur_val = gg + (x << l);
                    int ones = __builtin_popcountll(gg / 2) + (x & 1ll);
                    int zeros = l - ones;


                    if (x == 1) {
                        --ones;
                    }


                    ans += dp[diff + 60][ones][zeros];

//
                }
            }
        }
//        LOG(l, ans.val);
    }
    cout << ans.val << '\n';

    if (debug) {
        mint res = stupid(n);
        if (ans.val != res.val) {
            LOG(n);
            LOG(ans.val, res.val);
            exit(4);
        }
        LOG("test");
    }
}


int main() {
//    freopen("input.txt", "r", stdin);
//    freopen("output.txt", "w", stdout);

    ios_base::sync_with_stdio(0);
    cin.tie(0);

    for (int ones = 0; ones < max_n; ones++) {
        for (int zeros = 0; zeros + ones < max_n; zeros++) {
            for (int use_one = 0; use_one <= ones; use_one++) {
                for (int use_zero = 0; use_zero <= zeros; use_zero++) {
                    int diff = use_one - use_zero;
                    if (diff <= 60 && diff >= -60) {
                        dp[diff + 60][ones][zeros] += C(ones, use_one) * C(zeros, use_zero);
                    }
                }
            }
        }
    }
//    return 0;

    int t = 1;
    cin >> t;
    for (int test = 1; test <= t; test++) {
        test_case();
    }
}

/*
KoRoVa!
*/

C++