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

#include <algorithm>
#include <chrono>
#include <cmath>
#include <functional>
#include <iostream>
#include <iterator>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

typedef long long ll;
typedef vector<int> vi;
typedef vector<vi> vi2;
typedef vector<vi2> vi3;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<pii> vii;
typedef vector<ll> vll;
typedef vector<vll> vll2;
typedef vector<vll2> vll3;
typedef vector<bool> vb;

#define el '\\n'
#define rep(i, a, b) for (int i = a; i < b; i++)
#define repi(i, a, b) for (int i = a; i >= b; i--)
#define repeat(n) for (int i = 0; i < n; i++)
#define umap unordered_map
#define uset unordered_set
#define vec vector
#define loop(a) for (auto &x : a)
#define all(a) begin(a), end(a)
#define mp make_pair
#define pb push_back

template <typename T> void debug0(const T &value) { cerr << value << endl; }

template <typename T, typename... Args>
void debug0(const T &value, const Args &...args) {
  cerr << value << ' ';
  debug0(args...);
}

template <typename... Args>
void check0(const bool predicate, const Args &...args) {
  if (!predicate) {
    debug0(args...);
    throw logic_error("wtf");
  }
}

#ifdef LOCAL
#define debug(...) debug0(__VA_ARGS__)
#define check(...) check0(__VA_ARGS__)
#else
#define debug(...)
#define check(...)
#endif

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

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

  int n;
  cin >> n;
  vll a(n);
  loop(a) { cin >> x; }

  ll target = 0;
  rep(i, 2, n) { target ^= a[i]; }
  debug("target", target);

  const ll INF = 1e18;
  const ll NA = -INF - 10;
  const int MAX_BIT = 50;
  vll3 cache(MAX_BIT, vll2(2, vll(2, NA)));

  // aa = a_after_final & ((1LL << which) - 1)
  // return max_a_after_final
  function<ll(int, ll, ll)> dfs = [&](int which, ll aa, ll bb) {
    if (which >= MAX_BIT) {
      return -INF;
    }

    /* debug(which, aa, bb); */

    ll mask = (1LL << which) - 1;

    // make sure that aa ^ bb == target in the end
    if ((aa ^ bb) != (mask & target)) {
      return -INF;
    }

    // make sure that aa <= a[0]
    if (aa > a[0]) {
      return -INF;
    }

    // make sure that aa + bb == a[0] + a[1] in the end
    if (((aa + bb) & mask) != ((a[0] + a[1]) & mask) || aa + bb > a[0] + a[1]) {
      return -INF;
    }

    if (aa > 0 && aa <= a[0] && aa + bb == a[0] + a[1] && (aa ^ bb) == target) {
      debug("which", which, "aa", aa, "bb", bb);
      return 0LL;
    }

    /* debug("a[0] & mask", a[0] & mask); */
    /* debug("a[1] & mask", a[1] & mask); */
    int status_aa = aa > (a[0] & mask) ? 1 : 0; // borrow or not
    int status_bb = bb < (a[1] & mask) ? 1 : 0; // carry or not
    int sum_carry = aa + bb > ((a[0] + a[1]) & mask) ? 1 : 0;
    /* debug("sum_carry", sum_carry); */
    /* status_bb = sum_carry; */
    if (cache[which][status_aa][status_bb] >= 0) {
      /* debug("here"); */
      return cache[which][status_aa][status_bb];
    }

    ll ans = -INF;
    rep(a_bit, 0, 2) {
      rep(b_bit, 0, 2) {
        ll new_aa = aa + ((ll)a_bit << which);
        ll new_bb = bb + ((ll)b_bit << which);
        ll ans1 = dfs(which + 1, new_aa, new_bb);
        if (ans1 < 0) {
          continue;
        }
        ans = max(ans, ans1 + ((ll)a_bit << which));
      }
    }

    return cache[which][status_aa][status_bb] = ans;
  };

  ll max_first_after = dfs(0, 0, 0);
  debug("max_first_after", max_first_after);
  /* debug("cache[2,0,0]", cache[2][0][0]); */
  if (max_first_after <= 0) {
    cout << -1 << el;
  } else {
    cout << a[0] - max_first_after << el;
  }
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <chrono>
#include <cmath>
#include <functional>
#include <iostream>
#include <iterator>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

typedef long long ll;
typedef vector<int> vi;
typedef vector<vi> vi2;
typedef vector<vi2> vi3;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<pii> vii;
typedef vector<ll> vll;
typedef vector<vll> vll2;
typedef vector<vll2> vll3;
typedef vector<bool> vb;

#define el '\\n'
#define rep(i, a, b) for (int i = a; i < b; i++)
#define repi(i, a, b) for (int i = a; i >= b; i--)
#define repeat(n) for (int i = 0; i < n; i++)
#define umap unordered_map
#define uset unordered_set
#define vec vector
#define loop(a) for (auto &x : a)
#define all(a) a.begin(), a.end()
#define mp make_pair
#define pb push_back

template <typename T> void debug0(const T &value) { cerr << value << endl; }

template <typename T, typename... Args>
void debug0(const T &value, const Args &...args) {
  cerr << value << ' ';
  debug0(args...);
}

/* void check0(const bool predicate, string msg) { */
/*     if (!predicate) { */
/*         throw logic_error(msg); */
/*     } */
/* } */

template <typename... Args>
void check0(const bool predicate, const Args &...args) {
  if (!predicate) {
    debug0(args...);
    throw logic_error("wtf");
  }
}

#ifdef LOCAL
#define debug(...) debug0(__VA_ARGS__)
#define check(...) check0(__VA_ARGS__)
#else
#define debug(...)
#define check(...)
#endif

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

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

  int n;
  cin >> n;
  vll a(n);
  loop(a) { cin >> x; }

  ll target = 0;
  rep(i, 2, n) { target ^= a[i]; }
  debug("target=", target);

  int highest = 0;
  {
    ll max_ab = max(a[0], a[1]);
    while ((1LL << highest) <= max_ab) {
      highest++;
    }
    /* highest++; */
    debug("max_ab    :", max_ab);
    debug("1<<highest:", (1LL << highest));
  }
  debug("highest=", highest);

  const ll INF = 1e18;
  const ll NA = -INF - 10;
  debug(NA);
  vec<vec<vec<ll>>> cache(50, vec<vec<ll>>(2, vec<ll>(2, NA)));

  // aa = a_after_final & ((1LL << which) - 1)
  // return max_a_after_final
  function<ll(int, ll, ll)> dfs = [&](int which, ll aa, ll bb) {
    /* debug(which, aa, bb); */
    ll mask = (1LL << which) - 1;
    if ((aa ^ bb) != (mask & target)) {
      return -INF;
    }

    if (aa > a[0] || aa + bb > a[0] + a[1]) {
      return -INF;
    }

    if (((aa + bb) & mask) != ((a[0] + a[1]) & mask)) {
      return -INF;
    }

    /* ll a_sub = (a[0] & mask) - aa; */
    /* if (a_sub < 0) { */
    /*     a_sub += 1LL << which;  // borrow */
    /* } */
    /* check(a_sub >= 0, "a_sub=", a_sub); */
    /* check(a_sub < 1LL << which, "a_sub=", a_sub); */
    /* ll b_add = bb - (a[1] & mask); */
    /* if (b_add < 0) { */
    /*     b_add += 1LL << which;  // carry */
    /* } */
    /* check(b_add >= 0, "b_add=", b_add); */
    /* check(b_add < 1LL << which, "b_add=", b_add); */
    /* if (a_sub != b_add) { */
    /*     return -INF; */
    /* } */
    /* { */
    /*     ll delta = abs(a_sub - b_add); */
    /*     if (__builtin_popcountll(delta) > 1) { */
    /*         return -INF; */
    /*     } */
    /* } */

    if (which > highest) {
      if (aa > 0 && aa + bb == a[0] + a[1]) {
        debug("aa=", aa, "bb=", bb);
      }
      return aa > 0 && aa + bb == a[0] + a[1] ? 0 : -INF;
    }

    int status_aa = aa > (a[0] & mask) ? 1 : 0; // borrow or not
    int status_bb = bb < (a[1] & mask) ? 1 : 0; // carry or not
    /* ll &c = cache[which][status_aa][status_bb]; */
    if (cache[which][status_aa][status_bb] != NA) {
      /* return cache[which][status_aa][status_bb]; */
    }

    ll ans = -INF;
    rep(a_bit, 0, 2) {
      rep(b_bit, 0, 2) {
        ll new_aa = aa + ((ll)a_bit << which);
        ll new_bb = bb + ((ll)b_bit << which);
        ll ans1 = dfs(which + 1, new_aa, new_bb);
        if (ans1 < 0) {
          continue;
        }
        ans = max(ans, ans1 + ((ll)a_bit << which));
      }
    }
    return cache[which][status_aa][status_bb] = ans;
  };

  ll max_first_after = dfs(0, 0, 0);
  debug("max_first_after", max_first_after);
  if (max_first_after <= 0) {
    cout << -1 << el;
  } else {
    cout << a[0] - max_first_after << el;
  }
}