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) { */
    if (aa <= a[0] && aa + bb == a[0] + a[1] && (aa ^ bb) == target) {
      debug("[found ans]", "which", which, "aa", aa, "bb", bb, "ans", 0);
      return 0LL;
    }

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

    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));
      }
    }

    if (make_tuple(which, a_borrow, b_carry) == make_tuple(4, 0, 0)) {
      debug("which", which, "aa", aa, "bb", bb, "a_borrow", a_borrow, "b_carry",
            b_carry, "sum_carry", sum_carry, "sum_carry2", sum_carry2, "ans",
            ans);
    }

    if (cache[which][a_borrow][b_carry] != NA &&
        cache[which][a_borrow][b_carry] != ans) {
      debug("[cache change]", "which", which, "aa", aa, "bb", bb, "a_borrow",
            a_borrow, "b_carry", b_carry, "cache",
            cache[which][a_borrow][b_carry], "ans", ans);
    }

    return cache[which][a_borrow][b_carry] = ans;
  };

  debug("cache[4][0][0]", cache[4][0][0]);
  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) 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;
  map<tuple<int, ll, ll>, ll> cache;
  /* map<pair<int, pair<ll, ll>>, ll> cache; */

  // 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("[found ans]", "which", which, "aa", aa, "bb", bb, "ans", 0);
      return 0LL;
    }

    /* auto status = mp(which, mp(aa, bb)); */
    auto status = make_tuple(which, aa, bb);
    if (cache.count(status)) {
      return cache[status];
    }

    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[status] = 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;
  }
}