Tetris 3D help

Problems and Contests
#1

Problem

I’ve been dying on this problem for the last >2 hours.

What I tried

  • Quadtree (but then I realized that the update/query time complexity of quadtrees can go up to O(\text{maximum coordinate}))
  • 2D lazy segment tree (got TLE :frowning:)

My 2D lazy segment tree implementation (I think it’s readable without comments):

#include <bits/stdc++.h>
using namespace std;

const int inf = 1e9;
const int N = 1050;
const int M = N;
int n=1, m=1, q;

struct segtree1D {
    struct node {
        int max=0, lazy=0;
    } tr[2*M];

    void push(int i){
        tr[i].max = max(tr[i].max, tr[i].lazy);
        if(i<m){
            tr[2*i].lazy = max(tr[2*i].lazy, tr[i].lazy);
            tr[2*i+1].lazy = max(tr[2*i+1].lazy, tr[i].lazy);
        }
        tr[i].lazy = 0;
    }

    void upd(int l, int r, int i, int il, int ir, int x){
        push(i);
        if(ir < l || r < il)
            return;
        if(l <= il && ir <= r){
            tr[i].lazy = max(tr[i].lazy, x);
            push(i);
            return;
        }

        int im = (il+ir)/2;
        upd(l, r, 2*i, il, im, x);
        upd(l, r, 2*i+1, im+1, ir, x);
        tr[i].max = max(tr[i].max, x);
    }

    void upd(int l, int r, int x){
        assert(1 <= l && l <= r && r <= m);
        upd(l, r, 1, 1, n, x);
    }

    int qry(int l, int r, int i, int il, int ir){
        push(i);
        if(ir < l || r < il)
            return -inf;
        if(l <= il && ir <= r)
            return tr[i].max;

        int im = (il+ir)/2;
        return max(qry(l, r, 2*i, il, im), qry(l, r, 2*i+1, im+1, ir));
    }

    int qry(int l, int r){
        assert(1 <= l && l <= r && r <= m);
        return qry(l, r, 1, 1, n);
    }
};

struct segtree2D {
    struct node {
        segtree1D seg;
        vector<tuple<int, int, int>> lazy;
    } tr[2*N];

    void push(int i){
        for(const auto&[l, r, x]: tr[i].lazy) {
            tr[i].seg.upd(l, r, x);
            if (i < n) {
                tr[2 * i].lazy.emplace_back(l, r, x);
                tr[2 * i + 1].lazy.emplace_back(l, r, x);
            }
        }
        tr[i].lazy.clear();
    }

    void upd(int l, int r, int l2, int r2, int i, int il, int ir, int x){
        push(i);
        if(ir < l || r < il)
            return;
        if(l <= il && ir <= r){
            tr[i].lazy.emplace_back(l2, r2, x);
            push(i);
            return;
        }

        int im = (il+ir)/2;
        upd(l, r, l2, r2, 2*i, il, im, x);
        upd(l, r, l2, r2, 2*i+1, im+1, ir, x);
        tr[i].seg.upd(l2, r2, x);
    }

    void upd(int l, int r, int l2, int r2, int x){
        assert(1 <= l && l <= r && r <= n);
        upd(l, r, l2, r2, 1, 1, n, x);
    }

    int qry(int l, int r, int l2, int r2, int i, int il, int ir){
        push(i);
        if(ir < l || r < il)
            return -inf;
        if(l <= il && ir <= r)
            return tr[i].seg.qry(l2, r2);

        int im = (il+ir)/2;
        return max(qry(l, r, l2, r2, 2*i, il, im),
                   qry(l, r, l2, r2, 2*i+1, im+1, ir));
    }

    int qry(int l, int r, int l2, int r2){
        assert(1 <= l && l <= r && r <= n);
        return qry(l, r, l2, r2, 1, 1, n);
    }
} seg;

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);

    int n2, m2;
    cin >> n2 >> m2 >> q;
    while(n<n2)
        n *= 2;
    while(m<m2)
        m *= 2;

    while(q--){
        // d by s block with height w
        int d, s, w, x, y;
        cin >> d >> s >> w >> x >> y;
        int r1 = x+1, r2 = x+d;
        int c1 = y+1, c2 = y+s;
        int h = seg.qry(r1, r2, c1, c2);
        seg.upd(r1, r2, c1, c2, h+w);
    }

    cout << seg.qry(1, n, 1, m);
}

Question

What can I do to speed above the above code? Or, is there a better solution than a 2D lazy segtree?

#2

Never mind. I solved this problem. Arpa’s comment on CodeForces was helpful.