Problem p02958 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <array>
#include <bitset>
#include <cmath>
#include <complex>
#include <cstdio>
#include <deque>
#include <fstream>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORR(i, a, b) for (int i = a; i >= b; --i)
#define ALL(c) (c).begin(), (c).end()

typedef long long ll;
typedef vector<int> VI;
typedef vector<ll> VL;
typedef vector<long double> VD;
typedef vector<VI> VVI;
typedef vector<VL> VVL;
typedef vector<VD> VVD;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;

template <typename T> void chmin(T &a, T b) {
  if (a > b)
    a = b;
}
template <typename T> void chmax(T &a, T b) {
  if (a < b)
    a = b;
}

int in() {
  int x;
  scanf("%d", &x);
  return x;
}
ll lin() {
  ll x;
  scanf("%lld", &x);
  return x;
}

int main() {
  int n;
  cin >> n;
  vector<int> p(n);
  for (int i = 0; i < n; i++)
    cin >> p[i];

  int a = -1;
  int b = -1;
  bool flag = false;
  for (int i = 0; i < n - 1; i++)
    if (p[i] > p[i + 1]) {
      if (!flag) {
        a = i;
        flag = true;
      } else {
        b = i + 1;
        break;
      }
    }

  if (a == -1) {
    cout << "YES" << endl;
    return 0;
  }

  int tmp = p[b];
  p[b] = p[a];
  p[a] = tmp;

  for (int i = 0; i < n - 1; i++)
    if (p[i] > p[i + 1]) {
      cout << "NO" << endl;
      return 0;
    }

  cout << "YES" << endl;

  return 0;
}

Fixed submission:
#include <algorithm>
#include <array>
#include <bitset>
#include <cmath>
#include <complex>
#include <cstdio>
#include <deque>
#include <fstream>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORR(i, a, b) for (int i = a; i >= b; --i)
#define ALL(c) (c).begin(), (c).end()

typedef long long ll;
typedef vector<int> VI;
typedef vector<ll> VL;
typedef vector<long double> VD;
typedef vector<VI> VVI;
typedef vector<VL> VVL;
typedef vector<VD> VVD;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;

template <typename T> void chmin(T &a, T b) {
  if (a > b)
    a = b;
}
template <typename T> void chmax(T &a, T b) {
  if (a < b)
    a = b;
}

int in() {
  int x;
  scanf("%d", &x);
  return x;
}
ll lin() {
  ll x;
  scanf("%lld", &x);
  return x;
}

int main() {
  int n;
  cin >> n;
  vector<int> p(n);
  for (int i = 0; i < n; i++)
    cin >> p[i];

  int a = -1;
  int b = -1;
  bool flag = false;
  for (int i = 0; i < n - 1; i++)
    if (p[i] > p[i + 1]) {
      if (!flag) {
        a = i;
        flag = true;
      } else {
        b = i + 1;
        break;
      }
    }

  if (a == -1) {
    cout << "YES" << endl;
    return 0;
  } else if (b == -1) {
    cout << "NO" << endl;
    return 0;
  }

  int tmp = p[b];
  p[b] = p[a];
  p[a] = tmp;

  for (int i = 0; i < n - 1; i++)
    if (p[i] > p[i + 1]) {
      cout << "NO" << endl;
      return 0;
    }

  cout << "YES" << endl;

  return 0;
}