Problem p02691 (C++) — buggy submission, verdict: Time Limit Exceeded

/*

i<j

j-i=A[i]+A[j]

A[i]+i=-A[j]+j

1 1
4 2


i=1 A[i]=1
j=4 A[j]=2

A[i]+i=2
-A[j]+j=2

-A[i]+i=0
A[j]+j=6
*/
#include <bits/stdc++.h>
using namespace std;

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

  int N;
  scanf("%d", &N);

  vector<int> A(N + 1);     // A[0]は欠番
  vector<int> A_add(N + 1); // A[0]は欠番
  vector<int> A_sub(N + 1); // A[0]は欠番

  for (int i = 1; i <= N; i++) {
    scanf("%d", &A[i]);
    A_add[i] = A[i] + i;
    A_sub[i] = -A[i] + i;
  }

  long long result = 0;

  sort(A_add.begin(), A_add.end());
  sort(A_sub.begin(), A_sub.end());

  // for (int i = 1; i <= N; i++) {
  // 	cout << A[i] << " " << A_add[i] << " " << A_sub[i] << "\\n";
  // }
  int a = 1;
  int s = 1;
  long a_count = 0;
  long s_count = 0;
  while (true) {
    a_count = 1;
    // cout << "A_add=" << A_add[a] << "\\n";
    while (true) {
      if (A_add[a + 1] == A_add[a]) {
        a_count++;
        a++;
        if (a > N) {
          goto a_fin;
        }
      } else {
        break;
      }
    }
  a_fin:
    s_count = 0;
    while (true) {
      if (A_sub[s] == A_add[a]) {
        // cout << "A_sub[s] == A_add[a]\\n";
        s_count++;
        s++;
        if (s > N) {
          goto s_fin;
        }
      } else if (A_sub[s] < A_add[a]) {
        // cout << "A_sub[s] < A_add[a]\\n";
        s++;
        if (s > N) {
          goto s_fin;
        }
      } else {
        // cout << "A_sub[s] > A_add[a]\\n";
        break;
      }
    }
  s_fin:
    // cout << A_add[a] << " " << a_count << " " << s_count << "\\n";
    result += a_count * s_count;
    // cout << "result =" << result << "\\n";
    a++;
    if (a > N || s > N) {
      break;
    }
  }
  cout << result;
  return 0;
}

Error: TLE

Fixed submission:

/*

i<j

j-i=A[i]+A[j]

A[i]+i=-A[j]+j

1 1
4 2


i=1 A[i]=1
j=4 A[j]=2

A[i]+i=2
-A[j]+j=2

-A[i]+i=0
A[j]+j=6
*/
#include <bits/stdc++.h>
using namespace std;

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

  int N;
  scanf("%d", &N);

  vector<int> A(N + 1);     // A[0]は欠番
  vector<int> A_add(N + 1); // A[0]は欠番
  vector<int> A_sub(N + 1); // A[0]は欠番

  for (int i = 1; i <= N; i++) {
    scanf("%d", &A[i]);
    A_add[i] = A[i] + i;
    A_sub[i] = -A[i] + i;
  }

  int result = 0;
  for (int i = 1; i <= N; i++) {
    for (int j = i + 1; j <= N; j++) {
      if (A_add[i] == A_sub[j]) {
        result++;
      }
    }
  }
  cout << result;

  // for (int i = 1; i <= N; i++) {
  // 	cout << A[i] << " " << A_add[i] << " " << A_sub[i] << "\\n";
  // }

  // long long result = 0;

  // sort(A_add.begin(), A_add.end());
  // sort(A_sub.begin(), A_sub.end());

  // int a = 1;
  // int s = 1;
  // int a_count = 0;
  // int s_count = 0;
  // while (true) {
  // 	a_count = 1;
  // 	// cout << "A_add=" << A_add[a] << "\\n";
  // 	while (true) {
  // 		if (A_add[a + 1] == A_add[a]) {
  // 			a_count++;
  // 			a++;
  // 			if (a > N) {
  // 				goto a_fin;
  // 			}
  // 		} else {
  // 			break;
  // 		}
  // 	}
  // a_fin:
  // 	s_count = 0;
  // 	while (true) {
  // 		if (A_sub[s] == A_add[a]) {
  // 			// cout << "A_sub[s] == A_add[a]\\n";
  // 			s_count++;
  // 			s++;
  // 			if (s > N) {
  // 				goto s_fin;
  // 			}
  // 		} else if (A_sub[s] < A_add[a]) {
  // 			// cout << "A_sub[s] < A_add[a]\\n";
  // 			s++;
  // 			if (s > N) {
  // 				goto s_fin;
  // 			}
  // 		} else {
  // 			// cout << "A_sub[s] > A_add[a]\\n";
  // 			break;
  // 		}
  // 	}
  // s_fin:
  // 	result += a_count * s_count;
  // 	// cout << "result =" << result << "\\n";
  // 	a++;
  // 	if (a > N || s > N) {
  // 		break;
  // 	}
  // }
  // cout << result;
  // return 0;
}