TREE2SUM - Editorial

editorial
, , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: mathmodel
Tester: mexomerf
Editorialist: iceknight1093

DIFFICULTY:

simple

PREREQUISITES:

Binary search

PROBLEM:

Consider an infinite binary tree, where the children of node i are nodes 2i and 2i+1.

The score of a node u is defined to be the sum of labels of the last N nodes on the path from 1 to u (or all of them, if there are less than N).

Find any node u whose score is exactly equal to X, or claim that none exist.

EXPLANATION:

First, note that considering any nodes u such that u \gt X is pointless: simply starting at such a node would give us a sum that’s larger than X.

Next, note that for any starting node u, its score can be computed using brute force in O(\log u) time.
This is because the path from u to 1 is exactly formed by the nodes u, \left\lfloor \frac u 2 \right\rfloor, \left\lfloor \frac u 4 \right\rfloor, \left\lfloor \frac u 8 \right\rfloor, \ldots

So, we can simply bruteforce over the ancestors of u, till we either reach the root or have visited N nodes.
Note that since we’ve restricted ourselves to only having to think about u \leq X, and X \leq 10^{18}, this process will take at most \log_2{10^{18}} \leq 60 steps which is pretty fast.

Now, let f(u, N) denote the score of u with parameter N.
Observe that f is monotonic: that is, f(u, N) \lt f(u+1, N) for any integer u \geq 1.
This is because

f(u, N) = u + \left\lfloor \frac u 2 \right\rfloor + \left\lfloor \frac u 4 \right\rfloor + \left\lfloor \frac u 8 \right\rfloor + \ldots + \left\lfloor \frac{u}{2^{N-1}} \right\rfloor \\[20pt] f(u+1, N) = u+1 + \left\lfloor \frac {u+1} 2 \right\rfloor + \left\lfloor \frac {u+1} 4 \right\rfloor + \left\lfloor \frac {u+1} 8 \right\rfloor + \ldots + \left\lfloor \frac{u+1}{2^{N-1}} \right\rfloor

and it’s easy to see that each term of the first summation is not larger than the corresponding term of the second summation; with u \lt u+1 making the second summation strictly larger.

So, we have with us a monotonic function that can be computed in \mathcal{O}(\log X) time for a fixed parameter.
Our aim is to find a parameter u for which f(u, N) = X. This is now a standard application of binary search on the range [1, X].

TIME COMPLEXITY:

\mathcal{O}(\log^2 X) per testcase.

CODE:

Author's code (Python) 
from math import  *

def f(x,n):
  n=min(n,60) 
  s=0
  for i in range(n):
    s+=x//(1<<i)
  return s

for _ in range(int(input())):
  flag=0
  x,n=map(int,input().split())
  l,r=0,x
  while(l<=r):
    node=(l+r)//2
    if(f(node,n)==x):
      print(node)
      flag=1
      break
    elif(f(node,n)>x):
      r=node-1
    else:
      l=node+1
  if(flag==0):
    print(-1)
Tester's code (C++) 
#include <bits/stdc++.h>
using namespace std;
#define int long long
int cal(int n, int u){
    int sum = 0;
    while(u > 0 && n > 0){
        if(sum > LLONG_MAX - u){
            return LLONG_MAX;
        }
        sum += u;
        u /= 2;
        n--;
    }
    return sum;
}
int32_t main() {
	int t;
	cin>>t;
	while(t--){
	    int x, n;
	    cin>>x>>n;
	    int l = 1;
	    int r = 1e18;
	    int ans = -1;
	    while(l <= r){
	        int mid = (l + r) / 2;
	        int temp = cal(n, mid);
	        if(temp == x){
	            ans = mid;
	            break;
	        }
	        if(temp < x){
	            l = mid + 1;
	        }else{
	            r = mid - 1;
	        }
	    }
	    cout<<ans<<"\n";
	}
}
Editorialist's code (PyPy3) 
for _ in range(int(input())):
    x, n = map(int, input().split())
    
    def f(mid):
        sm, cur = 0, mid
        for i in range(n):
            sm += cur
            cur //= 2
            if cur == 0: break
        return sm
    
    lo, hi = 1, x
    while lo < hi:
        mid = (lo + hi) // 2
        
        if f(mid) < x: lo = mid + 1
        else: hi = mid
    if f(lo) == x: print(lo)
    else: print(-1)
1 Like