Submission #1236124

#TimeUsernameProblemLanguageResultExecution timeMemory
1236124hugo_pmMosaic (IOI24_mosaic)C++20
100 / 100
169 ms94064 KiB
#include "mosaic.h"
#include <cassert>
#include <vector>
#include <iostream>
#define int long long
using i32 = int32_t;
using namespace std;

struct Query {
    int r1, r2, c1, c2;
};

int prod(int u, int v) {
    return !(u || v);
}

struct MyVec {
    vector<int> vals, pref_sum, pref_idx_val;
    int n;
    MyVec() = default;
    MyVec(vector<int> _a) : vals(_a), n(_a.size()) {
        pref_sum.push_back(0);
        for (int x : vals)
            pref_sum.push_back(pref_sum.back() + x);
        pref_idx_val.push_back(0);
        for (int i = 0; i < (int)vals.size(); ++i)
            pref_idx_val.push_back(pref_idx_val.back() + i * vals[i]);
    }
    int sum(int l, int r) {
        if (l > r) return 0;
        assert(0 <= l && l <= r && r < n);
        return pref_sum[r+1] - pref_sum[l];
    }
    int sum_incr(int l, int r) {
        if (l > r) return 0;
        assert(0 <= l && l <= r && r < n);
        return pref_idx_val[r+1] - pref_idx_val[l] - (l-1) * sum(l, r);
    }
    int sum_decr(int l, int r) {
        if (l > r) return 0;
        assert(0 <= l && l <= r && r < n);
        int nb_vals = r-l+1;
        return (nb_vals + 1) * sum(l, r) - sum_incr(l, r);
    }
    int rectangle(int l, int r, int k) {
        assert(0 <= l && l <= r && r < n);
        return k*sum(l, r) + sum_incr(l - (k-1), l - 1) + sum_decr(r+1, r + (k-1));
    }
};

vector<int> solve(vector<int> X, vector<int> Y, vector<Query> queries) {
    int N = X.size();
    while (N < 4) {
        X.push_back(0), Y.push_back(0), ++N;
    }
    vector<vector<int>> grid(N, vector<int>(3));
    for (int row = 0; row < 3; ++row) {
        grid[row].resize(N);
    }
    grid[0] = X;
    for (int row = 0; row < N; ++row) {
        grid[row][0] = Y[row];
    }
    for (int row = 1; row < N; ++row) {
        for (int col = 1; col < (int)grid[row].size(); ++col) {
            grid[row][col] = prod(grid[row-1][col], grid[row][col-1]);
        }
    }
    vector<MyVec> sumRow(3), sumCol(3);
    for (int row = 0; row < 3; ++row) {
        sumRow[row] = MyVec(grid[row]);
    }
    for (int col = 0; col < 3; ++col) {
        vector<int> c;
        for (int row = 0; row < N; ++row) {
            c.push_back(grid[row][col]);
        }
        sumCol[col] = MyVec(c);
    }

    vector<int> border;
    for (int row = N-1; row >= 2; --row) {
        border.push_back(grid[row][2]);
    }
    for (int col = 3; col < N; ++col) { // don't push grid[2][2] twice
        border.push_back(grid[2][col]);
    }
    MyVec sumBorder(border);
    vector<int> answers;

    auto project = [&] (int row, int col) {
        return (N-3) - (row-col);
    };
    assert(project(N-1, 2) == 0);
    assert(project(2, N-1) == (int)border.size() - 1);

    for (auto [r1,r2,c1,c2] : queries) {
        int ans = 0;
        for (int row = r1; row <= min(r2, 2ll); ++row) {
            ans += sumRow[row].sum(c1, c2);
        }
        r1 = max(r1, 3ll);
        for (int col = c1; col <= min(c2, 2ll); ++col) {
            ans += sumCol[col].sum(r1, r2);
        }
        c1 = max(c1, 3ll);
        if (r1 <= r2 && c1 <= c2) {
            int delta1 = project(r1,c1), delta2 = project(r2, c2);
            int diagRect = min(c2-c1, r2-r1)+1;
            if (delta1 > delta2) swap(delta1, delta2);
            ans += sumBorder.rectangle(delta1, delta2, diagRect);
        }
        answers.push_back(ans);
    }
    return answers;
}

vector<long long> mosaic(vector<i32> _X, vector<i32> _Y,
vector<i32> T, vector<i32> B, vector<i32> L, vector<i32> R) {
    int Q = (int)T.size();
    vector<int> X(_X.begin(), _X.end()), Y(_Y.begin(), _Y.end());
    vector<Query> queries;
    for (int i = 0; i < Q; ++i) {
        queries.push_back({T[i], B[i], L[i], R[i]});
    }
    return solve(X, Y, queries);
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...