Segment Tree Codechef TLE

511 Views Asked by At

I am trying to solve this CodeChef problem:

There are N coins kept on the table, numbered from 0 to N - 1. Initially, each coin is kept tails up.
You have to perform two types of operations:

  1. Flip all coins numbered between A and B inclusive. This is represented by the command "0 A B"

  2. Answer how many coins numbered between A and B inclusive are heads up. This is represented by the command "1 A B".

Input: The first line contains two integers, N and Q. Each of the next Q lines are either of the form "0 A B" or "1 A B" as mentioned above.

Output: Output 1 line for each of the queries of the form "1 A B" containing the required answer for the corresponding query.

What I have used is a segment tree. So that every time user enter a query of type 1 A B the output is the sum at that interval [A,B]. However I am getting a Time Limit Exceeded error. I believe the error is due to the update step 0 A B. After updating the elements in the array I reconstruct the tree. The code is given below. Can someone help me with a faster way to update?

BTW - I am getting the desired output for the sample input.

public class SegmentTree
{
    private int[] tree;
    private int maxsize;
    private int height;
    private static int elems[];
    private  final int STARTINDEX = 0; 
    private  final int ENDINDEX;
    private  final int ROOT = 0;

    public SegmentTree(int size)
    {
        height = (int)(Math.ceil(Math.log(size) /  Math.log(2)));
        maxsize = 2 * (int) Math.pow(2, height) - 1;
        tree = new int[maxsize];
        ENDINDEX = size - 1; 
    }

    private int leftchild(int pos)
    {
        return 2 * pos + 1;
    }

    private int rightchild(int pos)
    {
        return 2 * pos + 2;
    }

    private int mid(int start, int end)
    {
        return (start + (end - start) / 2); 
    }

    private int getSumUtil(int startIndex, int endIndex, int queryStart, int queryEnd, int current)
    {
        if (queryStart <= startIndex && queryEnd >= endIndex)
        {
            return tree[current];
        }

        if (endIndex < queryStart || startIndex > queryEnd)
        {
            return 0;
        }

        int mid = mid(startIndex, endIndex);

        return  getSumUtil(startIndex, mid, queryStart, queryEnd, leftchild(current)) 
                 + getSumUtil( mid + 1, endIndex, queryStart, queryEnd, rightchild(current));
    }

    public int getSum(int queryStart, int queryEnd)
    {
        if(queryStart < 0 || queryEnd > tree.length)
        {
            return -1;
        }

        return getSumUtil(STARTINDEX, ENDINDEX, queryStart, queryEnd, ROOT);
    }

    private int constructSegmentTreeUtil(int startIndex, int endIndex, int current)
    {
        if (startIndex == endIndex)
        {
            tree[current] = elems[startIndex];
            return tree[current];   
        }

        int mid = mid(startIndex, endIndex);

        tree[current] = constructSegmentTreeUtil(startIndex, mid, leftchild(current))
                           + constructSegmentTreeUtil(mid + 1, endIndex, rightchild(current));

        return tree[current];
    }

    public void constructSegmentTree()
    {
        constructSegmentTreeUtil(STARTINDEX, ENDINDEX, ROOT);   
    }

    public static void main(String[]args) throws IOException
    {
        BufferedReader buf = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer str = new StringTokenizer(buf.readLine());
        int n = Integer.parseInt(str.nextToken());
        int q = Integer.parseInt(str.nextToken());
        SegmentTree segmentTree = new SegmentTree(n);
        int elements[] = new int[n];
        for(int i = 0; i < n; i++) {
            elements[i] = 0;
        }
        elems = elements;
        segmentTree.constructSegmentTree();
        while (q-- > 0) {
            str = new StringTokenizer(buf.readLine());
            int x = Integer.parseInt(str.nextToken());
            int a = Integer.parseInt(str.nextToken());
            int b = Integer.parseInt(str.nextToken());
            if(x == 0) {
                for(int j = a; j <= b; j++)
                {
                    elems[j] = elems[j]^1;
                }
                segmentTree.constructSegmentTree();
            }
            else {
                int num = segmentTree.getSum(a, b);
                System.out.println(num);
            }
        }
    }   
}

EDIT:

According to GeeksForGeeks, tree construction costs O(n) and the update method is O(log n). So here are the new methods for update:

private void updateTreeUtil(int startIndex, int endIndex, int updatePos, int update, int current)
{
    if ( updatePos < startIndex || updatePos > endIndex)
    {
        return;
    }

    tree[current] = tree[current] + update;

    if (startIndex != endIndex)
    {
        int mid = mid(startIndex, endIndex);
        updateTreeUtil(startIndex, mid, updatePos, update, leftchild(current));
        updateTreeUtil(mid+1, endIndex, updatePos, update, rightchild(current));
    }
}

public void update(int update, int updatePos)
{
    int updatediff = update - elems[updatePos];
    elems[updatePos] = update;
    updateTreeUtil(STARTINDEX, ENDINDEX, updatePos, updatediff, ROOT);
}

And now the if loop in main method modified to this:

if(x == 0) {
    for(int j = a; j <= b; j++)
    {
        segmentTree.update(elems[j]^1, j);
    }
}

But still getting TLE error.

1

There are 1 best solutions below

2
On

In the tutorial of GeeksForGeeks, their running time of update is O(log n), in case of updating a single element. However, when doing update for an interval, you have to use Lazy Propagation to ensure O(log n) update time, which is basically only update nodes which are visited, and hence ensure the sum of visited nodes are correct. You may search for many good tutorial on Lazy Propagation, for example:

http://se7so.blogspot.hk/2012/12/segment-trees-and-lazy-propagation.html

Wish that helps.