Codeforces Round #719 Guess the K-th Zero (Hard version) - Segment tree

Solution

Base on the easy version, we still binary search in the hard version. So, we query the sum of $[1, mid]$.
Now the problem is, after each time, we need change the element from 0 to 1.

The direct thinking is using a map to cache all the queries. And every time, we update all the values in the map base on the result. But of course, it will get TLE.

So, we use a segment tree to cache and update the result of queries. We have 2 update operations and 1 search operation:

1. For the query to the system. We update the $mid$'s value. And mark $mid$ as queried. Also, it is mandatory to ignore and clear lazy operations in this position.
2. $+1$ for all cached queries between $[ans, n]$.
3. Search the cache of $mid$.

Implementation details:

The node of segment tree:

static class Node {
    int left, right, value = -1;
    int plus = 0;
    boolean flag = false;
}

• left, right represent the range of current node.
• value is the cached result of query. If left != right, then the value has no meaning.
• plus is the lazy operation. Which represent the value which need to add on all cached result for current range.
• flag represent whether this pos was queried or not. If left != right, then it has no meaning.

Push down the lazy operation:

private static void push_down(int p) {
    if (nodes[p].plus != 0) {
        nodes[p * 2].plus += nodes[p].plus;
        nodes[p * 2 + 1].plus += nodes[p].plus;
        nodes[p].plus = 0;
    }
}

For the three operations of segment tree, should note that:

1. For update mid's value. Because it is the latest result from system. So, the plus should be ignore. Just force change the value.
2. For search mid's value. If the flag is false, then return -1 (means no cache found) no matter what plus is.

Code: