Time Limit Exceeded with Segment Tree data structure

Question

I am trying to solve this problem

Chef Ceil has some matchsticks in his kitchen.

Detail of matchsticks:

There are N matchsticks in total. They are numbered from to 0 to N-1 inclusive. All matchsticks have same length. But they may have different rates of burning. For ith matchstick, we denote bi as the time required for that matchstick to completely burn-down if lighted at one end. The matchsticks have uniform rate of burning. If lighted at both ends simultaneously, the matchstick will take only half of the original time to burn down.

Arrangement:

He ties rear end of the all the matchsticks together at one point and the front end is kept free. The matchstick numbered i is adjacent to matchstick numbered i+1 for all 0<= i <=N-2. Bodies of matchsticks do not touch each other, except at rear end. All matchsticks are kept on the floor. Task: There are Q queries, in each query we ask: If he lights the free end of all matchsticks numbered between L and R both inclusive, what will be the time needed for all matchsticks to get completely burnt.

Input:

First line of input contains one integer N, total number of matchsticks. The next line contains N space separated integers bi, where bi is the time required for ith matchstick to completely burn-down if lighted at one end. Next line contains one integer Q, total number of queries you need to answer. Then there are Q queries in next Q lines. Each line has two space separated integers L and R.

Output:

Print Q lines, one for each query, printing the answer for each query, that is, the time it will take for all the matchsticks to get completely burn down. Every time you must print your answer with 1 decimal place.

Constraints:

\$1 <= N <= 105\$
\$1<= bi <= 108\$
\$1 <= Q <= 105\$
\$0 <= L <= R <= N-1\$

Examples:

Input - 1 5 1 0 0

Output - 5.0

Input - 2 3 5 1 0 1

Output - 4.0

Input - 18 3 4 2 1 5 7 9 7 10 5 12 3 1 1 2 1 3 2 1 4 10

Output - 9.0

Here is my solution to this problem:

#include <cstdio>
#include <vector>
#include <cmath>

#define rep(i, n) for(int i = 0; i < n; ++i)
#define RANGE_MIN 1
#define RANGE_MAX 2

using namespace std;

typedef vector<int> vi;
typedef vector<int> tl;
vi segment_tree; //tree for finding minimum
vi segment_tree_max; //tree for finding maximum
tl input;

//initialization of the segment trees
void init_segment_tree(int N){
    int length = (int)(2 * pow(2.0, floor((log((double)N) / log(2.0)) + 1)));
    segment_tree.resize(length,0);
    segment_tree_max.resize(length,0);
}


void build_segment_tree(int code,tl A,int node, int b, int e){
    if(b==e){
        if(code==RANGE_MIN) segment_tree[node] = b;
        else segment_tree_max[node] = b;
    } else {
        int leftIdx = 2*node;
        int rightIdx = 2*node + 1;
        build_segment_tree(code,A,leftIdx,b,(b+e)/2);
        build_segment_tree(code,A,rightIdx,(b+e)/2+1,e);
        if(code == RANGE_MIN){
            int lContent = segment_tree[leftIdx];
            int rContent = segment_tree[rightIdx];
            int lValue = A[lContent];
            int rValue = A[rContent];
            segment_tree[node] = (lValue>rValue) ? rContent : lContent;
        } else {
            int lContent = segment_tree_max[leftIdx];
            int rContent = segment_tree_max[rightIdx];
            int lValue = A[lContent];
            int rValue = A[rContent];
            segment_tree_max[node] = (lValue>rValue) ? lContent : rContent;
        }
    }
}


int query(int code,tl A, int node, int b,int e,int i,int j){
if(i>e || j<b) return -1;
if(b>=i && e<=j && code == RANGE_MIN) return segment_tree[node];
if(b>=i && e<=j && code == RANGE_MAX) return segment_tree_max[node];

int p1 = query(code,A,2*node,b,(b+e)/2,i,j);
int p2 = query(code,A,2*node + 1, (b+e)/2+1,e,i,j);
if(p1==-1) return p2;
if(p2==-1) return p1;
if(code == RANGE_MIN) return (A[p1]<=A[p2]) ? p1 : p2;
if(code == RANGE_MAX) return (A[p1]<=A[p2]) ? p2 : p1;
}

int main(){
    int n,L,R,Q;
    int time;
    int min,max1,max3;
    float max,max2;
    scanf("%d",&n);
    rep(i,n){
        scanf("%d",&time);
        input.push_back(time);
    }
    init_segment_tree(n);
    build_segment_tree(1,input,1,0,n-1);
    build_segment_tree(2,input,1,0,n-1);
    scanf("%d",&Q);
    while(Q--){
        scanf("%d%d",&L,&R);
        min = input[query(1,input,1,0,n-1,L,R)]; //find the minimum in [L,R]
        max1 = max3 =0;
        if (L>0) max1 = input[query(2,input,1,0,n-1,0,L-1)]; //find the maximum in [0,L-1]
        max2 = (float)(input[query(2,input,1,0,n-1,L,R)]-min)/2; //find maximum in [L,R] which is burning twice as fast
        if (R<n-1) max3 = input[query(2,input,1,0,n-1,R+1,n-1)]; //find maximum in [R+1,n-1]
        max = max1;
        if(max2>max) max = max2;
        if(max3>max) max = max3;
        printf("%0.1f\n",min+max);
    }
}

I don't know what's wrong with it, it gets TLE. I implemented the segment tree data as follows:

• \$O(2N)\$ for building 2 segment trees, one to find the minimum and another to find the maximum.
• \$O(4Qlog(N))\$ since I run 4 Range_Min(Max)_querys for each query \$Q\$.

Everything seems to be working but I still don"t understand why I am getting TLE.

Answer 1

There are some optimizations that may help

1. You are passing input vector by assignment (as tl A) to query() and build_segment_tree(). By doing this you are essentially copying the whole vector (of size \$O(10^5)\$ in worst case) in each recursive call for those functions.

   Your query() is of linear \$O(n)\$ complexity and build_segment_tree() is of quadratic \$O(N^2)\$ complexity. Pass it by const reference instead (as const tl &A).

2. For segment_tree, you have constructed it with indices instead of an actual value, so querying and constructing segment tree has an overhead of dereferencing (A[i] is *(A+i)) and possibly relatively more cache misses (depending on system cache size) as well since you are trying to access two vectors alternatively by doing this.

Answer 2

You have written your code in a way that makes it extremely difficult to understand, and almost impossible without the problem statement nearby. What is Q? Well, we have to return to the statement of the problem to find out it is the number of queries. Likewise, you use typedefs and macros that obfuscate the meaning of your code, and you use globals for reasons I don't understand.

After some study, I have determined that your algorithm is as follows:

1. Read in the parameters of the problem.
2. Construct a segment tree that supports queries of the form: What is the burn time of the shortest match stick in the range [i, j]?
3. Construct a segment tree that supports queries of the form: What is the burn time of the longest match stick in the range [i, j]?
4. Read in a query's parameters L, R.
5. Determine the burn time of the shortest match stick in the range [L, R] (min).
6. Determine the burn time of the longest match stick in the range [0, L-1] (max1).
7. Determine the burn time of the longest match stick in the range [L, R] (max2 -- this requires special treatment as it burns twice as fast after min time has elapsed).
8. Determine the burn time of the longest match stick in the range [R+1, N] (max3).
9. Print min + max(max1, max2, max3).
10. Repeat 4-9 while queries remain.

Let me offer an alternate version of your code that follows the same algorithm, with commentary interspersed:

#include <algorithm>
#include <cmath>
#include <functional>
#include <iostream>
#include <iterator>
#include <utility>
#include <vector>

namespace chef_ceil {

I have chosen to use iostream instead of cstdio; we're writing C++, not C. Also, I've omitted using namespace std; (see the top two answers of this SO question, and all of your macros, typedefs, and global variables.

I have also put everything into a namespace.

/**
 * A data structure that can answer queries of the form "what is the minimum
 * value between indexes L and R?" With Comp=std::greater<T>, finds maxima instead.
 *
 * Requires O(N) storage for an input with N values.
 */
template <typename T, typename Comp=std::less<T>>
class SegmentTree {
  public:
    /**
     * Constructs a SegmentTree over the given values.
     *
     * Requires O(N) time to construct.
     */
    static SegmentTree New(const std::vector<T>& ts, Comp comp=Comp());

    /**
     * Find the smallest (largest) value that has an index between l and r inclusive.
     *
     * Requires O(log N) time.
     */
    T query(std::size_t l, std::size_t r) const;

    std::size_t num_values() const { return num_values_; }

  private:
    SegmentTree(std::vector<T>&& nodes, Comp comp, std::size_t num_values)
      : nodes_(std::move(nodes)), comp_(comp), num_values_(num_values) {}

    T query_helper(std::size_t l, std::size_t r, std::size_t node,
                   std::size_t b, std::size_t e) const;

    std::vector<T> nodes_;
    Comp comp_;
    std::size_t num_values_;

};

C++ is an object-oriented language; it is sensible to define a class that encapsulates behavior instead of throwing the behavior of this data structure into various parts of the code.

namespace {
std::size_t num_nodes(std::size_t num_leaves) {
  // Round up to the nearest power of two to get the number of leaves in the
  // complete binary tree, then multiply by 2 (normally you would subtract 1,
  // but we leave node 0 empty and use node 1 as the root).
  return std::lrint(std::exp2(1 + std::ceil(std::log2(num_leaves))));
}

Note the comment explaining why I'm using such a complicated formula.

template <typename T, typename Comp>
T build_from_node(std::vector<T>& nodes, const std::vector<T>& ts, Comp comp,
                  std::size_t node, std::size_t l, std::size_t r) {
  if (l == r) {
    nodes[node] = ts[l];
    return ts[l];
  }
  int left_child = 2 * node;
  int right_child = 2 * node + 1;

  const T left_value = build_from_node(nodes, ts, comp, left_child, l, (l + r) / 2);
  const T right_value = build_from_node(nodes, ts, comp, right_child, (l + r) / 2 + 1, r);
  nodes[node] = comp(left_value, right_value) ? left_value : right_value;
  return nodes[node];
}
}  // namespace

template <typename T, typename Comp>
SegmentTree<T, Comp> SegmentTree<T, Comp>::New(const std::vector<T>& ts, Comp comp) {
  std::vector<T> nodes;
  nodes.reserve(num_nodes(ts.size()));
  build_from_node(nodes, ts, comp, 1, 0, ts.size() - 1);

  return SegmentTree(std::move(nodes), comp, ts.size());
}

template <typename T, typename Comp>
T SegmentTree<T, Comp>::query(std::size_t l, std::size_t r) const {
  if (l > r || l < 0 || r >= num_values_) throw std::out_of_range("bad query");

  return query_helper(l, r, 1, 0, num_values_ - 1);
}

template <typename T, typename Comp>
T SegmentTree<T, Comp>::query_helper(std::size_t l, std::size_t r, std::size_t node,
                                     std::size_t b, std::size_t e) const {
  if (l <= b && e <= r) return nodes_[node];

  const auto left_child = 2 * node;
  const auto right_child = 2 * node + 1;
  const auto mid = (b + e) / 2;

  if (l > mid)  // left child is irrelevant
    return query_helper(l, r, right_child, mid + 1, e);
  if (r <= mid)  // right child is irrelevant
    return query_helper(l, r, left_child, b, mid);

  const auto left_value = query_helper(l, r, left_child, b, mid);
  const auto right_value = query_helper(l, r, right_child, mid + 1, e);
  return comp_(left_value, right_value) ? left_value : right_value;
}

Note that the use of a comparison function as a template parameter means that I don't have to know which data structure I'm querying (max or min), and that's one less magic constant running around.

namespace {
std::vector<int> read_match_lengths(std::istream& is) {
  int num_matches;
  is >> num_matches;
  std::vector<int> match_lengths;
  match_lengths.reserve(num_matches);
  std::copy_n(std::istream_iterator<int>(is), num_matches, std::back_inserter(match_lengths));
  return match_lengths;
}

int compute_burn_time_for_range(const SegmentTree<int>& min_tree,
                                const SegmentTree<int, std::greater<int>>& max_tree,
                                int l, int r) {
  const int light_time = min_tree.query(l, r);
  const int left_burn_time = l > 0 ? max_tree.query(0, l - 1) : 0;
  const int right_burn_time = r + 1 < max_tree.num_values()
                              ? max_tree.query(r + 1, max_tree.num_values() - 1)
                              : 0;
  // The amount of time it takes to burn the longest match in the lit range
  // *after* all the match rears are lit is (b_i - light_time) / 2.
  const int lit_burn_time = (max_tree.query(l, r) - light_time) / 2;

  return light_time + std::max({left_burn_time, right_burn_time, lit_burn_time});
}
}   // namespace
}  // namespace chef_ceil

Here I've moved a bunch of code out of main; this makes it easier to test, in addition to being easier to read.

int main(int, char**) {
  const auto match_lengths = chef_ceil::read_match_lengths(std::cin);
  int num_queries;
  std::cin >> num_queries;

  const auto min_tree = chef_ceil::SegmentTree<int>::New(match_lengths);
  const auto max_tree = chef_ceil::SegmentTree<int, std::greater<int>>::New(match_lengths);

  for (int i = 0; i < num_queries; ++i) {
    int l, r;
    std::cin >> l >> r;
    std::cout << chef_ceil::compute_burn_time_for_range(min_tree, max_tree, l, r) << "\n";
  }
}

Note: I have only done minimal testing of the code in this example.