TIOJ Infor Online Judge Problem 1827. Yet another simple task

TIOJ1827.cpp

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

struct Node {
    int sum;
    Node *lc, *rc;
    Node(int s);
    Node(Node*, Node*);
};

Node::Node(int s) : sum(s), lc(nullptr), rc(nullptr) {}
Node::Node(Node *l, Node *r) : sum(l->sum + r->sum), lc(l), rc(r) {}

Node* create(int l, int r) {
    if(l + 1 == r)
        return new Node(0);
    else
        return new Node(create(l, (l + r) / 2), create((l + r) / 2, r));
}

int sum(Node *n, int r, int nl, int nr) {
    if(nr <= r)
        return n->sum;
    else if(r > nl)
        return sum(n->lc, r, nl, (nl + nr) / 2) +
               sum(n->rc, r, (nl + nr) / 2, nr);
    else
        return 0;
}

Node* increase(Node *n, int x, int nl, int nr) {
    if(nl + 1 == nr)
        return new Node(n->sum + 1);
    else if(x < (nl + nr) / 2)
        return new Node(increase(n->lc, x, nl, (nl + nr) / 2), n->rc);
    else
        return new Node(n->lc, increase(n->rc, x, (nl + nr) / 2, nr));
}

Node* nodes[100001];

bool ok(int n, int p, int k, int s) {
    return sum(nodes[min(p + s + 1, n)], s, 0, n) -
           sum(nodes[max(p - s    , 0)], s, 0, n) >= k;
}

int main(){
    int n, m;
    scanf("%d %d", &n, &m);
    nodes[0] = create(0, n);
    for(int i = 0; i < n; i++) {
        int b;
        scanf("%d", &b);
        b--;
        nodes[i + 1] = increase(nodes[i], b, 0, n);
    }
    for(int i = 0; i < m; i++) {
        int p, k;
        scanf("%d %d", &p, &k);
        int l = 0, r = n;
        while(l < r) {
            int s = (l + r) / 2;
            ok(n, p, k, s) ? (r = s) : (l = s + 1);
        }
        printf("%d\n", r);
    }
}

TIOJ1827.hs

import Control.Monad

data Node =
      Leaf   Int           -- sum
    | Branch Int Node Node -- sum, lchild, rchild

middle :: Integral a => a -> a -> a
middle = ((`div` 2) .) . (+)

create :: Int -> Int -> Node
create l r
    | l + 1 == r = Leaf 0
    | otherwise  = Branch 0 (create l m) (create m r)
    where m = middle l r
    
sumof :: Node -> Int -> Int -> Int -> Int
sumof (Leaf v) r _ nr
    | nr <= r   = v
    | otherwise = 0
sumof (Branch v lc rc) r nl nr
    | nr <= r   = v
    | r  > nl   = (sumof lc r nl m) + (sumof rc r m nr)
    | otherwise = 0
    where m = middle nl nr

increase :: Node -> Int -> Int -> Int -> Node
increase (Leaf v) _ _ _ = Leaf (v + 1)
increase (Branch v lc rc) x nl nr
    | x < m     = Branch (v + 1) (increase lc x nl m) rc
    | otherwise = Branch (v + 1) lc (increase rc x m nr)
    where m = middle nl nr

tonodes0 :: Int -> [Int] -> [Node]
tonodes0 n (h:t) = increase (head $ nodes) h 0 n : nodes where nodes = tonodes0 n t
tonodes0 n _ = [create 0 n]

tonodes :: Int -> [Int] -> [Node]
tonodes = (reverse .) . (. reverse) . tonodes0

ok :: [Node] -> Int -> Int -> Int -> Int -> Bool
ok nodes n p k s = (sumof (nodes !! min (p + s + 1) n) s 0 n) - (sumof (nodes !! max (p - s) 0) s 0 n) >= k

bsearch :: (Int -> Bool) -> Int -> Int -> Int
bsearch condition l r
    | l == r      = r
    | condition m = bsearch condition l m
    | otherwise   = bsearch condition (m+1) r
    where m = middle l r

query :: Int -> [Node] -> IO ()
query n nodes = do
    [p, k] <- fmap (map read . words) getLine
    print $ bsearch (ok nodes n p k) 0 n

main :: IO ()
main = do
    [n, m] <- fmap (map read . words) getLine
    nodes <- fmap (tonodes n . map (+(-1)) . map read . words) getLine
    replicateM m $ query n nodes
    return ()