Problem p03225 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define all(v) (v).begin(), (v).end()
#define sz(v) ((int)((v).size()))
typedef long long ll;
typedef pair<int, int> ii;

const int MAXN = 305;
const int dx[4] = {0, 1, 0, -1};
const int dy[4] = {1, 0, -1, 0};

int n, m, diagCount;

int d1[2 * MAXN][2 * MAXN];
int d2[2 * MAXN][2 * MAXN];
int p1[2 * MAXN][2 * MAXN];
int p2[2 * MAXN][2 * MAXN];
char a[MAXN][MAXN];

inline bool ok(int x, int y) { return (0 <= x && x < n) && (0 <= y && y < m); }

inline void makePrefixSum(int *arr) {
  for (int i = 1; i < diagCount; i++) {
    arr[i] += arr[i - 1];
  }
}

inline int prefSum(int *p, int a, int b, int dist) {
  if (a > b) {
    int t = a;
    a = b;
    b = t;
  }
  if (b >= diagCount) {
    b = diagCount - 1;
  }
  if (a < 0) {
    a = 0;
  }

  b -= dist;
  if (a > b)
    return 0;

  int res = p[b];
  if (a > 0)
    res -= p[a - 1];
  return res;
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cout.precision(15);
  cout.setf(ios::fixed);

  cin >> n >> m;
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }
  diagCount = n + m - 1;

  for (int i = 0; i < n; i++) {
    for (int j = 0; j < m; j++) {
      if (a[i][j] == '#') {
        int z1 = i + j;
        int z2 = i + (m - j - 1);
        d1[z1][z2]++;
        d2[z2][z1]++;
      }
    }
  }

  ll ans = 0;
  for (int dist = 2; dist <= n + m - 2; dist += 2) {
    for (int z1 = 0; z1 < diagCount; z1++) {
      for (int z2 = 0; z2 < diagCount; z2++) {
        p1[z1][z2] = 0;
        p2[z1][z2] = 0;
      }
    }
    for (int z1 = 0; z1 < diagCount; z1++) {
      for (int z2 = 0; z2 < diagCount; z2++) {
        if (d1[z1][z2] == 1 && z2 + dist < diagCount &&
            d1[z1][z2 + dist] == 1) {
          p1[z1][z2] = 1;
        }
        if (d2[z2][z1] == 1 && z1 + dist < diagCount &&
            d2[z2][z1 + dist] == 1) {
          p2[z2][z1] = 1;
        }
      }
    }
    for (int z = 0; z < diagCount; z++)
      makePrefixSum(p1[z]);
    for (int z = 0; z < diagCount; z++)
      makePrefixSum(p2[z]);

    ll curDistAns = 0;
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < m; j++) {
        if (a[i][j] == '#') {
          int cellAns = 0;
          for (int k = 0; k < 4; k++) {
            int xs = i + dist * dx[k];
            int ys = j + dist * dy[k];
            int xf = i + dist * dx[(k + 1) & 3];
            int yf = j + dist * dy[(k + 1) & 3];
            int z1s = xs + ys;
            int z2s = xs + (m - ys - 1);
            int z1f = xf + yf;
            int z2f = xf + (m - yf - 1);
            if (z1s == z1f) {
              if (0 <= z1s && z1s < diagCount) {
                int add = prefSum(p1[z1s], z2s, z2f, dist);
                cellAns += add;
              }
            } else if (z2s == z2f) {
              if (0 <= z2s && z2s < diagCount) {
                int add = prefSum(p2[z2s], z1s, z1f, dist);
                cellAns += add;
              }
            } else {
              throw 0;
            }
            int xm = (xs + xf) / 2;
            int ym = (ys + yf) / 2;
            if (ok(xs, ys) && ok(xm, ym)) {
              if (a[xs][ys] == '#' && a[xm][ym] == '#') {
                cellAns--;
              }
            }
          }
          curDistAns += cellAns;
        }
      }
    }
    ans += curDistAns;
  }

  cout << ans << endl;
}


Fixed submission:
// #include <bits/stdc++.h>
#include <cstdio>
#include <iostream>
using namespace std;
#define all(v) (v).begin(), (v).end()
#define sz(v) ((int)((v).size()))
typedef long long ll;
typedef pair<int, int> ii;

const int MAXN = 305;
const int dx[4] = {0, 1, 0, -1};
const int dy[4] = {1, 0, -1, 0};

int n, m, diagCount;

int d1[2 * MAXN][2 * MAXN];
int d2[2 * MAXN][2 * MAXN];
int p1[2 * MAXN][2 * MAXN];
int p2[2 * MAXN][2 * MAXN];
char a[MAXN][MAXN];

inline bool ok(int x, int y) { return (0 <= x && x < n) && (0 <= y && y < m); }

inline void makePrefixSum(int *arr) {
  for (int i = 1; i < diagCount; i++) {
    arr[i] += arr[i - 1];
  }
}

inline int prefSum(int *p, int a, int b, int dist) {
  if (a > b) {
    int t = a;
    a = b;
    b = t;
  }
  if (b >= diagCount) {
    b = diagCount - 1;
  }
  if (a < 0) {
    a = 0;
  }

  b -= dist;
  if (a > b)
    return 0;

  int res = p[b];
  if (a > 0)
    res -= p[a - 1];
  return res;
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cout.precision(15);
  cout.setf(ios::fixed);

  cin >> n >> m;
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }
  diagCount = n + m - 1;

  for (int i = 0; i < n; i++) {
    for (int j = 0; j < m; j++) {
      if (a[i][j] == '#') {
        int z1 = i + j;
        int z2 = i + (m - j - 1);
        d1[z1][z2]++;
        d2[z2][z1]++;
      }
    }
  }

  ll ans = 0;
  for (int dist = 2; dist <= n + m - 2; dist += 2) {
    for (int z1 = 0; z1 < diagCount; z1++) {
      for (int z2 = 0; z2 < diagCount; z2++) {
        p1[z1][z2] = 0;
        p2[z1][z2] = 0;
      }
    }
    for (int z1 = 0; z1 < diagCount; z1++) {
      for (int z2 = 0; z2 < diagCount; z2++) {
        if (d1[z1][z2] == 1 && z2 + dist < diagCount &&
            d1[z1][z2 + dist] == 1) {
          p1[z1][z2] = 1;
        }
        if (d2[z2][z1] == 1 && z1 + dist < diagCount &&
            d2[z2][z1 + dist] == 1) {
          p2[z2][z1] = 1;
        }
      }
    }
    for (int z = 0; z < diagCount; z++)
      makePrefixSum(p1[z]);
    for (int z = 0; z < diagCount; z++)
      makePrefixSum(p2[z]);

    ll curDistAns = 0;
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < m; j++) {
        if (a[i][j] == '#') {
          int cellAns = 0;
          for (int k = 0; k < 4; k++) {
            int xs = i + dist * dx[k];
            int ys = j + dist * dy[k];
            int xf = i + dist * dx[(k + 1) & 3];
            int yf = j + dist * dy[(k + 1) & 3];
            int z1s = xs + ys;
            int z2s = xs + (m - ys - 1);
            int z1f = xf + yf;
            int z2f = xf + (m - yf - 1);
            if (z1s == z1f) {
              if (0 <= z1s && z1s < diagCount) {
                int add = prefSum(p1[z1s], z2s, z2f, dist);
                cellAns += add;
              }
            } else if (z2s == z2f) {
              if (0 <= z2s && z2s < diagCount) {
                int add = prefSum(p2[z2s], z1s, z1f, dist);
                cellAns += add;
              }
            } else {
              throw 0;
            }
            int xm = (xs + xf) / 2;
            int ym = (ys + yf) / 2;
            if (ok(xs, ys) && ok(xm, ym)) {
              if (a[xs][ys] == '#' && a[xm][ym] == '#') {
                cellAns--;
              }
            }
          }
          curDistAns += cellAns;
        }
      }
    }
    ans += curDistAns;
  }

  cout << ans << endl;
}