HackerRank Quadrant Queries problem solution

In this HackerRank Quadrant Queries problem solution, we have given n points on a plane and each point is described by (x,y) coordinates and we need to perform some queries given by the user.

Problem solution in Python.

import bisect
from math import sqrt
import optparse
import sys
from random import randint as rand

_Rx = [3, 2, 1, 0]
_Ry = [1, 0, 3, 2]

def create_test(N, Q, fp):
    r = [(rand(-2, 2), rand(-2, 2)) for i in range(N)]
    print(N, file = fp)
    for ri in r:
        print('%d %d' % ri, file = fp)
    ops = [ "X", "Y", "C" ]
    queries = [(ops[rand(0, 2)], rand(1, N), rand(1, N)) for q in range(Q)]
    print(Q, file = fp)
    for op, i, j in queries:
        if i > j: 
            i, j = j, i
        print('%1s %d %d' % (op, i, j), file = fp)
    return [r, queries]
    
def quadrant(x, y):
    '''
    Return the quadrant the point r is in.
    '''
    if x > 0:
        if y > 0: return 0
        else: return 3
    else:
        if y > 0: return 1
        else: return 2
    
def getInput(f, shift = -1):
    '''
    The first line contains N, the number of points. N lines follow.
    The ith line contains xi and yi separated by a space.  The next line
    contains Q the number of queries. The next Q lines contain one query
    each, of one of the above forms.  All indices are 1 indexed.

    Return the points r and the queries q, in [r, q]
    '''
    l = f.readline().strip()
    n = int(l)
    r = []
    for i in range(n):
        l = f.readline().strip()
        x, y = l.split()
        r.append([int(x), int(y)])
    l = f.readline().strip()
    Q = int(l)
    queries = []
    for k in range(Q):
        l = f.readline().strip()
        op, i, j = l.split()
        queries.append((op, int(i) + shift, int(j) + shift))
    return [r, queries]

class QuadrantBook:
    '''
    A data structure to facilitate the processing of quadrant queries.
    '''
    def __init__(self, r, factor = 32):
        N = len(r)
        self.blocksize = int(sqrt(factor * N))
        nblocks = (N - 1) // self.blocksize + 1
        # self.inq[b][q] is the list of points in the quadrant q with indices
        # between b L and (b+1)L, where L is the blocksize.
        self.inq = [[[] for q in range(4) ] for b in range(nblocks)]
        for i, ri in enumerate(r):
            b = self.getBlock(i)
            q = quadrant(ri[0], ri[1])
            self.inq[b][q].append(i)

    def countInQuadrants( self, i, j):
        '''
        Count the number of points in each quadrant with indices between i and j,
        inclusive.
        '''
        bi = self.getBlock(i)
        bj = self.getBlock(j)
        cq = [0] * 4
        if bi == bj:
            for q, idx in enumerate(self.inq[bi]):
                ki, kj = getIndex(idx, i, j)
                cq[q] = kj - ki
            return cq
        else:
            for q in range(4):
                ni = self.inq[bi][q]
                nj = self.inq[bj][q]
                ki = bisect.bisect_left(ni, i)
                kj = bisect.bisect_right(nj, j)
                cq[q] = len(ni) - ki 
                for nb in self.inq[bi + 1: bj]:
                    cq[q] += len(nb[q])
                cq[q] += kj
            return cq

    def reflect( self, i, j, axis):
        '''
        Update the list of points in each quadrant after reflection along the
        given axis.
        '''
        if axis == "X":
            R = _Rx
            W = [(0, 3), (1, 2)]
        else:
            R = _Ry
            W = [(0, 1), (2, 3)]    
        bi = self.getBlock( i)
        bj = self.getBlock( j)
        if bi == bj:
            self.reflectInBlock( i, j, bi, R)
        else:
            ni = self.inq[bi]
            kiq = [bisect.bisect_left(idx, i) for idx in ni]
            rfxi = [ni[q][ki:] for q, ki in enumerate(kiq)]
            for q, ki in enumerate(kiq):
                ni[q] = ni[q][:ki] + rfxi[R[q]]   
            nj = self.inq[bj]
            kjq = [bisect.bisect_right( idx, j) for idx in nj]
            rfxj = [nj[q][:kj] for q, kj in enumerate(kjq)]
            for q, kj in enumerate( kjq):
                nj[q] = rfxj[R[q]] + nj[q][kj:]    
            for nb in self.inq[bi + 1: bj]:
                for w1, w2 in W:
                    nb[w1], nb[w2] = nb[w2], nb[w1]
          
    def getBlock( self, i ):
        '''
        Get the block index for point i.
        '''
        return i // self.blocksize


    def reflectInBlock(self, i, j, b, R):
        '''
        Update the list of points in each quadrant after reflection along the
        given axis, assuming that i and j are in the same block b.
        '''
        nb = self.inq[b]
        kij = [getIndex( idx, i, j) for idx in nb]
        rf = [nb[q][ki: kj] for q, (ki, kj) in enumerate(kij)]
        for q, (ki, kj) in enumerate(kij):
            nb[q] = nb[q][:ki] + rf[R[q]] + nb[q][kj:]

def getIndex( lst, i, j):
    '''
    In the sorted list, find ki such that lst[ki-1]<i<=lst[ki], and kj such
    that lst[kj]<=j<lst[kj+1]. If i<lst[0], ki=0; if j>lst[-1], kj=len(lst).
    This convention splices the list at the right places.

    @param i<=j
    '''
    ki = bisect.bisect_left(lst, i)
    kj = bisect.bisect_right(lst, j, ki)
    return [ki, kj]

def processQueries(r, queries, factor = 32):
    qbook = QuadrantBook(r, factor)
    for op, i, j in queries:
        if op.upper() == 'C':
            cq = qbook.countInQuadrants(i, j)
            print('%d %d %d %d' % tuple(cq), file = sys.stdout)
        elif op.upper() in ['X', 'Y']:
            qbook.reflect( i, j, op)


if __name__ == '__main__':
    usage = '%prog'
    opt = optparse.OptionParser(usage)
    opt.add_option('-N', type = 'int', default = 100)
    opt.add_option('-Q', type = 'int', default = 10)
    opt.add_option('-c', '--create', action = 'store_true', default = False, help = 'Create test case instead of running the query.')
    opt.add_option('-i', '--input', type='string', default = None, help = 'Input file.')
    opt.add_option('-b', '--block-factor', type = 'float', default = 32., help = 'Use sqrt(b * N) as the block size.')
    opts, args = opt.parse_args()
    if opts.create:
        r, queries = create_test(opts.N, opts.Q, sys.stdout)
    else:
        r, queries = getInput(opts.input and open(opts.input, 'r') or sys.stdin)
        processQueries(r, queries, opts.block_factor)
  

Problem solution in Java.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;

public class Solution {
    private static final int[] BC = new int[1<<16];
    public static void main(String[] argv) throws Exception {
        prep();
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String line1 = br.readLine();
        int N = Integer.parseInt(line1);
        ArrayList<String> dump = new ArrayList<String>(100000);
        int[] x = new int[N/32+1];
        int[] y = new int[N/32+1];
        long rs = System.currentTimeMillis();
        for(int i=0;i<N;i++){
            String line2 = br.readLine();
            String[] tmp2 = line2.split(" ");
            if(tmp2[0].charAt(0) != '-') x[i/32] |= 1<< i%32;
            if(tmp2[1].charAt(0) != '-') y[i/32] |= 1<< i%32;
        }
        long re = System.currentTimeMillis();
        String line3 = br.readLine();
        int Q = Integer.parseInt(line3);
        boolean wasC = false;
        int[] tmp1 = new int[N/32+1];
        int[] tmp24= new int[N/32+1];
        int[] tmp2 = new int[N/32+1];
        int[] cnt1 = new int[N/32+1];
        int[] cnt24= new int[N/32+1];
        int[] cnt2 = new int[N/32+1];
        long ws = System.currentTimeMillis();
        for(int i=0;i<Q;i++){
            String line4 = br.readLine();
            String[] tmp4 = line4.split(" ");
            int sindex = Integer.parseInt(tmp4[1]);
            int eindex = Integer.parseInt(tmp4[2]);
            int sm = (sindex-1)%32;
            int sn = (sindex-1)/32;
            int em = eindex%32;
            int en = eindex/32;
            int es = eindex-sindex+1;
            if(tmp4[0].equals("X")) {
                if(sn == en){
                    int mask = ((1<<(es))-1)<<(sm);
                    y[en] ^= mask;
                }else{
                    int mask = -1;
                    int fmask = mask<<(sm);
                    int emask = (1<<(em))-1;
                    y[sn] ^= fmask;
                    for(int j=sn+1;j<en;j++) y[j] ^= mask;
                    y[en] ^= emask;
                }
                wasC = false;
            }else if(tmp4[0].equals("Y")) {
                if(sn == en){
                    int mask = ((1<<(es))-1)<<(sm);
                    x[en] ^= mask;
                }else{
                    int mask = -1;
                    int fmask = mask<<(sm);
                    int emask = (1<<(em))-1;
                    x[sn] ^= fmask;
                    for(int j=sn+1;j<en;j++) x[j] ^= mask;
                    x[en] ^= emask;
                }
                wasC = false;
            }else{
                /*
                if(!wasC || wasC){
                    for(int j=0;j<x.length;j++) {
                        tmp1[j]  = x[j] & y[j];
                        tmp24[j] = x[j] ^ y[j];
                        tmp2[j]  = tmp24[j] & y[j];
                        cnt1[j] = bitCount(tmp1[j]);
                        cnt24[j]= bitCount(tmp24[j]);
                        cnt2[j] = bitCount(tmp2[j]);
                    }
                }
                */
                int maskes = ((1<<(es))-1)<<(sm);
                int maskall = -1;
                int fmask = maskall<<(sm);
                int emask = (1<<(em))-1;
                // 1st quadrant x bit: 1, y bit: 1 (x & y)
                int c1 = 0;
                if(sn == en){
                    c1 += bitCount(x[en] & y[en] & maskes);
                }else{
                    c1 += bitCount(x[sn] & y[sn] & fmask);
                    for(int j=sn+1;j<en;j++) c1 += bitCount(x[j] & y[j]);
                    c1 += bitCount(x[en] & y[en] & emask);
                }
                // 2nd quadrant x bit: 0, y bit: 1
                // 4th quadrant x bit: 1, y bit: 0
                // x xor y = c2 + c4
                // (x xor y) & y = c2
                int c24 = 0;
                int c2 = 0;
                if(sn == en){
                    int t2 = (x[en] ^ y[en]) & maskes;
                    c24 += bitCount(t2);
                    c2 += bitCount(t2 & y[en]);
                }else{
                    int t2 = (x[sn] ^ y[sn]) & fmask;
                    c24 += bitCount(t2);
                    c2 += bitCount(t2 & y[sn]);
                    for(int j=sn+1;j<en;j++) {
                        t2 = x[j] ^ y[j];
                        c24 += bitCount(t2);
                        c2 += bitCount(t2 & y[j]);
                    }
                    t2 = (x[en] ^ y[en]) & emask;
                    c24 += bitCount(t2);
                    c2 += bitCount(t2 & y[en]);
                }
                int c4 = c24 - c2;
                // 3rd quadrant x bit: 0, y bit: 0 (total - c1 - c2 - c4)
                int c3 = eindex - sindex + 1 - c1 - c2 - c4;
                dump.add(c1+" "+c2+" "+c3+" "+c4);
                //System.out.println(c1+" "+c2+" "+c3+" "+c4);
                wasC = true;
            }
        }
        long we = System.currentTimeMillis();
        for(int i=0;i<dump.size();i++) System.out.println(dump.get(i));
        br.close();
        //System.out.println("R:"+(re-rs));
        //System.out.println("W:"+(we-ws));
    }
    private static void prep() {
        int max = 1<<16;
        int step1 = 1<<8;
        for(int i=0;i<step1;i++) BC[i] = Integer.bitCount(i);
        for(int i=step1;i<max;i++){
            BC[i] = BC[i&0xFF] + BC[(i>>8)&0xFF];
        }
    }
    
    private static int bitCount(int i){
        return BC[i&0xFFFF] + BC[(i>>16)&0xFFFF];
    }
}

Problem solution in C++.

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
using namespace std;

const int maxN = 100000;

const int R[4][4] =
{
  {0, 1, 2, 3},
  {3, 2, 1, 0},
  {1, 0, 3, 2},
  {2, 3, 0, 1}
};

int N;
int px[maxN], py[maxN];

int Q;

struct Node
{
  int l, r;
  int nq[4]; // number of points in each quadrant, applying reflect mask on children nodes (exclude self)
  int mask; // reflect mask X(0) Y(1)
  Node *lt, *rt;
};

Node *root;

Node *bst_build (int l, int r)
{
  Node *p = (Node *)malloc(sizeof(Node));
  p->l = l, p->r = r;
  memset(p->nq, 0, sizeof(p->nq));
  p->mask = 0;
  p->lt = p->rt = NULL;

  if (r - l > 1)
  {
    int m = (l + r) / 2;
    Node *lt = bst_build(l, m);
    Node *rt = bst_build(m, r);
    p->lt = lt;
    p->rt = rt;
  }

  if (r - l > 1)
  {
    for (int i = 0; i < 4; ++i)
      p->nq[i] = p->lt->nq[i] + p->rt->nq[i];
  }
  else
  {
    int q;
    if (px[l] > 0)
      if (py[l] > 0)
        q = 0;
      else
        q = 3;
    else
      if (py[l] > 0)
        q = 1;
      else
        q = 2;
    p->nq[q] = 1;
  }

  return p;
}

void bst_mask (int ax, Node *p, int l, int r)
{
  if (p->l == l && p->r == r)
  {
    p->mask ^= 1 << ax;
    return;
  }

  if (l < p->lt->r)
    bst_mask(ax, p->lt, l, min(r, p->lt->r));
  if (r > p->rt->l)
    bst_mask(ax, p->rt, max(l, p->rt->l), r);
  
  for (int i = 0; i < 4; ++i)
  {
    p->nq[i] = p->lt->nq[R[p->lt->mask][i]] + p->rt->nq[R[p->rt->mask][i]];
  }
}

void bst_query (Node *p, int l, int r, int mask, int res[4])
{
  mask ^= p->mask;
  
  if (p->l == l && p->r == r)
  {
    for (int i = 0; i < 4; ++i)
    {
      res[i] += p->nq[R[mask][i]];
    }
    return;
  }

  if (l < p->lt->r)
    bst_query(p->lt, l, min(r, p->lt->r), mask, res);
  if (r > p->rt->l)
    bst_query(p->rt, max(l, p->rt->l), r, mask, res);
}

int main ()
{
  scanf("%d", &N);
  for (int i = 0; i < N; ++i)
    scanf("%d%d", &px[i], &py[i]);

  root = bst_build(0, N);

  scanf("%d", &Q);
  for (int i = 0; i < Q; ++i)
  {
    char c;
    int l, r;
    scanf(" %c%d%d", &c, &l, &r);
    --l;
    switch (c)
    {
      case 'X':
      {
        bst_mask(0, root, l, r);
        break;
      }
      case 'Y':
      {
        bst_mask(1, root, l, r);
        break;
      }
      case 'C':
      {
        int res[4];
        memset(res, 0, sizeof(res));
        bst_query(root, l, r, 0, res);
        printf("%d %d %d %dn", res[0], res[1], res[2], res[3]);
        break;
      }
    }
  }
}

Problem solution in C.

#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
 
enum {
  XAXIS = 0,
  YAXIS = 1
};
 

 
typedef struct node node;
struct node {
  int count;
  node *left, *right;
};
 
node *pool;
int nnodes, npool;
 
typedef struct tree tree;
struct tree {
  node *root[4];
  int nlevels;
};
 
node *newnode() {
  assert(nnodes < npool);
  return &pool[nnodes++];
}
 
node *mktree(int level) {
  node *np;
 
  np = newnode();
  if (level == 0)
    return np;
 
  np->left = mktree(level-1);
  np->right = mktree(level-1);
  return np;
}
 
tree *inittree(int n) {
  tree *tp;
  int i;
 
  tp = (tree *) malloc(sizeof(*tp));
  if (tp == NULL)
    exit(1);
  
  tp->nlevels = 0;
  while ((1 << tp->nlevels) < n)
    tp->nlevels++;
 
  npool = 4 * (1 << (tp->nlevels+1));
  pool = calloc(npool, sizeof(node));
  if (pool == NULL)
    exit(1);
 
  for (i = 0; i < 4; i++)
    tp->root[i] = mktree(tp->nlevels);
 
  return tp;
}
 
void _add(node *np, int lvl, int idx, int x) {
  int lt, rt, mid;
 
  np->count++;
 
  if (lvl == 0)
    return;
 
  lt = (1<<lvl) * idx;
  mid = lt + (1<<(lvl-1));
  rt = lt + (1<<lvl) - 1;
 
  assert(lt <= x && x <= rt);
  if (x < mid)
    _add(np->left, lvl-1, 2*idx, x);
  else
    _add(np->right, lvl-1, 2*idx + 1, x);
}
 
/* Add a value at x in quadrant q. */
void add(tree *tp, int x, int q) {
  _add(tp->root[q], tp->nlevels, 0, x);
}
 
int _count(node *np, int lvl, int idx, int l, int r) {
  int lt, rt;
 
  lt = (1<<lvl) * idx;
  rt = lt + (1<<lvl) - 1;
 
  /* Out of bounds. */
  if (r < lt) return 0;
  if (l > rt) return 0;
 
  /* Whole range. */
  if (l <= lt && rt <= r)
    return np->count;
 
  /* Left and right children. */
  return _count(np->left, lvl-1, 2*idx, l, r) +
    _count(np->right, lvl-1, 2*idx + 1, l, r);
}
 
/* count: print # in each quadrant from l -> r inclusive. */
void count(tree *tp, int l, int r) {
  int q[4];
  int i;
 
  for (i = 0; i < 4; i++)
    q[i] = _count(tp->root[i], tp->nlevels, 0, l, r);
  printf("%d %d %d %dn", q[0], q[1], q[2], q[3]);
}
 
void _flip(node **np, int lvl, int idx, int l, int r, int typ) {
  int lt, rt, i;
  node *tmp, *lp[4], *rp[4];
 
  lt = (1<<lvl) * idx;
  rt = lt + (1<<lvl) - 1;
 
  /* Out of bounds. */
  if (r < lt) return;
  if (l > rt) return;
 
  /* Whole range. */
  if (l <= lt && rt <= r) {
    assert(typ == XAXIS || typ == YAXIS);
    switch (typ) {
    case XAXIS:
      tmp = np[0]; np[0] = np[3]; np[3] = tmp;
      tmp = np[1]; np[1] = np[2]; np[2] = tmp;
      break;
    case YAXIS:
      tmp = np[0]; np[0] = np[1]; np[1] = tmp;
      tmp = np[2]; np[2] = np[3]; np[3] = tmp;
      break;
    default:
      assert(0);
      break;
    }
    return;
  }
 
  assert(lvl > 0);
 
  /* Left and right children. */
  for (i = 0; i < 4; i++) {
    lp[i] = np[i]->left;
    rp[i] = np[i]->right;
  }
 
  _flip(lp, lvl-1, 2*idx, l, r, typ);
  _flip(rp, lvl-1, 2*idx + 1, l, r, typ);
 
  for (i = 0; i < 4; i++) {
    np[i]->left = lp[i];
    np[i]->right = rp[i];
    np[i]->count = lp[i]->count + rp[i]->count;
  }
}
 
/* flip: flip points about axis. */
void flip(tree *tp, int l, int r, int typ) {
  _flip(tp->root, tp->nlevels, 0, l, r, typ);
}
 
int quadrant(int x, int y) {
  if (x > 0 && y > 0) return 0;
  if (x < 0 && y > 0) return 1;
  if (x < 0 && y < 0) return 2;
  if (x > 0 && y < 0) return 3;
  /* shouldn't get here */
  assert(0);
  return -1;
}
 
int main() {
  int i, x, y, a, b, npoints, nqueries;
  char query;
  tree *tp;
 
  if (scanf(" %d", &npoints) != 1)
    return 1;
  tp = inittree(npoints);
  for (i = 0; i < npoints; i++) {
    if (scanf(" %d %d", &x, &y) != 2)
      return 1;
    if (x == 0 || y == 0)
      return 1;
    add(tp, i, quadrant(x, y));
  }
 
  if (scanf(" %d", &nqueries) != 1)
    return 1;
  for (i = 0; i < nqueries; i++) {
    if (scanf(" %c %d %d", &query, &a, &b) != 3)
      return 1;
    switch (query) {
    case 'C':
      count(tp, a-1, b-1);
      break;
    case 'X':
      flip(tp, a-1, b-1, XAXIS);
      break;
    case 'Y':
      flip(tp, a-1, b-1, YAXIS);
      break;
    default:
      return 1;
    }
  }
  
  return 0;
}