Problem p03157 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <cstdio>
#include <vector>

typedef long long llong;

constexpr int kMaxN = 400;
constexpr int kMaxNN = kMaxN * kMaxN;

char maze[kMaxN + 3][kMaxN + 3];
int union_set_id[kMaxNN + 3];
int union_set_count0[kMaxNN + 3];
int union_set_count1[kMaxNN + 3];
bool tmp[kMaxNN + 3];

int GetRoot(int i) {
  if (union_set_id[i] != i) {
    return GetRoot(union_set_id[i]);
  }
  return i;
}

void UnionSetUnion(int i, int j) {
  i = GetRoot(i);
  j = GetRoot(j);
  if (i == j) {
    return;
  }
  union_set_id[i] = j;
  union_set_count0[j] += union_set_count0[i];
  union_set_count1[j] += union_set_count1[i];
}

int main() {
  int n, m;
  scanf("%d %d", &n, &m);

  for (int i = 0; i < n; ++i) {
    scanf("%s", maze[i]);
    for (int j = 0; j < m; ++j) {
      union_set_id[i * m + j] = i * m + j;
      union_set_count0[i * m + j] = maze[i][j] == '.';
      union_set_count1[i * m + j] = maze[i][j] == '#';
    }
  }

  int goi[] = {1, -1, 0, 0};
  int goj[] = {0, 0, 1, -1};

  for (int i = 0; i < n; ++i) {
    for (int j = 0; j < m; ++j) {
      char c0 = maze[i][j];
      for (int k = 0; k < 4; ++k) {
        int i1 = i + goi[k];
        int j1 = j + goj[k];
        if (0 <= i1 && i1 < n && 0 <= j1 && j1 < m) {
          char c1 = maze[i1][j1];
          if (c0 != c1) {
            // printf("union (%d,%d) and (%d,%d)\n", i, j, i1, j1);
            UnionSetUnion(i * m + j, i1 * m + j1);
          }
        }
      }
    }
  }

  int nm = n * m;
  for (int i = 0; i < nm; ++i) {
    tmp[i] = false;
  }

  llong total = 0;
  for (int i = 0; i < nm; ++i) {
    int r = GetRoot(i);
    if (!tmp[r]) {
      tmp[r] = true;
      total += llong(union_set_count0[r]) * llong(union_set_count1[r]);
    }
  }

  printf("%lld\n", total);
  return 0;
}

Fixed submission:
#include <cstdio>
#include <vector>

typedef long long llong;

constexpr int kMaxN = 400;
constexpr int kMaxNN = kMaxN * kMaxN;

char maze[kMaxN + 3][kMaxN + 3];
int union_set_id[kMaxNN + 3];
int union_set_count0[kMaxNN + 3];
int union_set_count1[kMaxNN + 3];
bool tmp[kMaxNN + 3];

int GetRoot(int i) {
  if (union_set_id[i] != i) {
    return union_set_id[i] = GetRoot(union_set_id[i]);
  }
  return i;
}

void UnionSetUnion(int i, int j) {
  i = GetRoot(i);
  j = GetRoot(j);
  if (i == j) {
    return;
  }
  union_set_id[i] = j;
  union_set_count0[j] += union_set_count0[i];
  union_set_count1[j] += union_set_count1[i];
}

int main() {
  int n, m;
  scanf("%d %d", &n, &m);

  for (int i = 0; i < n; ++i) {
    scanf("%s", maze[i]);
    for (int j = 0; j < m; ++j) {
      union_set_id[i * m + j] = i * m + j;
      union_set_count0[i * m + j] = maze[i][j] == '.';
      union_set_count1[i * m + j] = maze[i][j] == '#';
    }
  }

  int goi[] = {1, -1, 0, 0};
  int goj[] = {0, 0, 1, -1};

  for (int i = 0; i < n; ++i) {
    for (int j = 0; j < m; ++j) {
      char c0 = maze[i][j];
      for (int k = 0; k < 4; ++k) {
        int i1 = i + goi[k];
        int j1 = j + goj[k];
        if (0 <= i1 && i1 < n && 0 <= j1 && j1 < m) {
          char c1 = maze[i1][j1];
          if (c0 != c1) {
            // printf("union (%d,%d) and (%d,%d)\n", i, j, i1, j1);
            UnionSetUnion(i * m + j, i1 * m + j1);
          }
        }
      }
    }
  }

  int nm = n * m;
  for (int i = 0; i < nm; ++i) {
    tmp[i] = false;
  }

  llong total = 0;
  for (int i = 0; i < nm; ++i) {
    int r = GetRoot(i);
    if (!tmp[r]) {
      tmp[r] = true;
      total += llong(union_set_count0[r]) * llong(union_set_count1[r]);
    }
  }

  printf("%lld\n", total);
  return 0;
}