HackerEarth Tom and Jerry love matrices solution

In this HackerEarth Tom and jerry love matrices problem solution Tom and Jerry are playing with matrices. Jerry gifts Tom with a matrix of size M x N. The matrix element have a unique property that each cell of the matrix has value Aij = X + i + j. Since Tom loves challenges, Jerry will give Tom Q queries. Each of these queries can be of any of the following three types:

1. R C1 C2 : Delete the cells in the Row  starting from column C1 and ending at column C2 (1 <= C1 <= C2 < N, 1 <= R <= M)
2. C R1 R2 : Delete the cells in column  starting from row R1 and ending at row R2 (1 <= R1 <= R2 <= M, 1 <= C <= N).
3. K: Take all the remaining elements in the array and sort the elements by their values. Now, return the Kth smallest element. If K is greater than the number of elements remaining, then print -1.(1 <= K <= 10^9)

You are required to print the output only in case of the queries of type 3. You must help Tom to solve the tasks and answer the queries given by Jerry.

HackerEarth Tom and Jerry love matrices problem solution.

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
 
using namespace __gnu_pbds;
using namespace std;
 
typedef long long ll;
int n, m, x, q;
 
ll total_count(int v)
{
    if (n > m)
        swap(n, m);
    ll ans = 0;
    if (2 <= v && v <= n + 1)
        ans += v * 1LL * (v - 1) / 2;
    else
        ans += n * 1LL * (n + 1) / 2;
    if (v >= n + 1)
        ans += min(m - n - 1, v - (n + 1)) * 1LL * n;
    if (m + 1 <= v && v <= m + n)
        ans += n * 1LL * (n + 1) / 2 - (m + n - v) * 1LL * (m + n - v + 1) / 2;
    return ans;
}
 
struct Vertex
{
    int left, right;
    ll value, lazy;
    Vertex *left_child = nullptr, *right_child = nullptr;
 
    Vertex(int lb, int rb)
    {
        value = 0;
        lazy = 0;
        left = lb;
        right = rb;
    }
 
    void extend()
    {
        if (!left_child)
        {
            int t = ((ll)left + right) >> 1;
            left_child = new Vertex(left, t);
            right_child = new Vertex(t + 1, right);
        }
    }
 
    void add(ll l, ll r)
    {
        if (right < l or r < left)
            return;
        else if (l == left && r == right)
        {
            value += r - l + 1LL;
            lazy += r - l + 1LL;
        }
        else
        {
            extend();
            int mid = ((ll)left + right) >> 1;
 
            if (lazy)
            {
                ll v = lazy / (right - left + 1);
                left_child->value += (mid - left + 1) * v;
                left_child->lazy += (mid - left + 1) * v;
                right_child->value += (right - mid) * v;
                right_child->lazy += (right - mid) * v;
                lazy = 0;
            }
 
            left_child->add(l, min((ll)mid, r));
            right_child->add(max(mid + 1LL, l), r);
            value = left_child->value + right_child->value;
        }
    }
 
    int get(int k, ll s)
    {
        if (left == right)
            return left;
 
        extend();
        int mid = ((ll)left + right) >> 1;
 
        if (lazy)
        {
            ll v = lazy / (right - left + 1);
            left_child->value += (mid - left + 1) * v;
            left_child->lazy += (mid - left + 1) * v;
            right_child->value += (right - mid) * v;
            right_child->lazy += (right - mid) * v;
            lazy = 0;
        }
 
        ll tl = total_count(left_child->right) - left_child->value - s;
 
        if (tl >= k)
            return left_child->get(k, s);
        else
            return right_child->get(k, s + left_child->value);
    }
 
    void dfs()
    {
        cout << "[ " << left << " " << right << " :: " << value << " " << lazy << " ]" << endl;
        if (left_child)
            left_child->dfs();
 
        if (right_child)
            right_child->dfs();
    }
};
 
signed main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
 
    cin >> n >> m >> x >> q;
 
    Vertex *root = new Vertex(2, n + m);
 
    while (q--)
    {
        int t;
        cin >> t;
 
        if (t <= 2)
        {
            int a, b, c;
            cin >> a >> b >> c;
            root->add(a + b, a + c);
        }
        else
        {
            int k;
            cin >> k;
            cout << (total_count(n + m) - root->value < k ? -1 : x + (ll)root->get(k, 0)) << '\n';
        }
    }
}

Second solution

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
const int MAX_N = 2e9 + 14;

struct Node {
    Node *L, *R;
    ll iman;
    int sina;

    Node() {
        L = R = 0;
        iman = sina = 0;
    }

    void arpa() {
        if (L) return;
        L = new Node();
        R = new Node();
    }

    void majid(int s, int e, int x, int l = 0, int r = MAX_N) {
        if (s <= l && r <= e) {
            sina += x;
            return;
        }
        if (e <= l || r <= s)
            return;
        arpa();
        int mid = (ll) l + r >> 1;
        L->majid(s, e, x, l, mid);
        R->majid(s, e, x, mid, r);
        iman = L->iman + L->sina * (ll) (mid - l) + R->iman + R->sina * (ll) (r - mid);
    }

    ll hamid(int s, int e, int l = 0, int r = MAX_N) {
        if (s <= l && r <= e) {
            return iman + sina * (ll) (r - l);
        }
        if (e <= l || r <= s) return 0;
        int mid = (ll) l + r >> 1;
        return (L ? L->hamid(s, e, l, mid) : 0) + (R ? R->hamid(s, e, mid, r) : 0) +
               sina * (ll) (min(r, e) - max(l, s));
    }
} root;

int n, m, x, q;

ll cnt(int i) {
    ll ans = -root.hamid(0, i + 1);
    if (i < min(n, m))
        ans += ll(i + 1) * (i + 2) / 2;
    else if (i < max(n, m))
        ans += ll(min(n, m) + 1) * (min(n, m)) / 2 + ll(i - min(n, m) + 1) * min(n, m);
    else
        ans += (ll) n * m - ll(n + m - i - 2) * (n + m - i - 1) / 2;
    return ans;
}

int main() {
    ios::sync_with_stdio(0), cin.tie(0);
    cin >> n >> m >> x >> q;
    x += 2;
    while (q--) {
        int t;
        cin >> t;
        if (t == 3) {
            int k;
            cin >> k;
            int lo = -1, hi = n + m - 1;
            while (hi - lo > 1) {
                int mid = (lo + hi) / 2;
                (cnt(mid) >= k ? hi : lo) = mid;
            }
            cout << (hi == n + m - 1 ? -1 : x + hi) << '\n';
        }
        else {
            int a, b, c;
            cin >> a >> b >> c;
            a--, b--;
            root.majid(a + b, a + c, 1);
        }
    }
}