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

#include <algorithm>
#include <array>
#include <bitset>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdint.h>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;
using ll = long long;
using ull = unsigned long long;
using uint = unsigned int;
using ldouble = long double;

// BDD,ZDD,kdtree, bdtree,bicalc, bolonoy, doloney, tree,
// chinesemod,segmenttree,daikusutora, saidairyuu, 2bugurahu, heirokenshutu,
// topologicalsort, kyourenketuseibun

#define REP(i, a, b) for (int i = a; i < b; ++i)
#define PRI(s) cout << s << endl
#define PRIY PRI("Yes")
#define PRIN PRI("No")

vector<vector<array<ll, 4>>> dp;
vector<vector<ll>> table;
int R, C, K;

int main() {
  cin >> R >> C >> K;
  dp.resize(R);
  table.resize(R);
  REP(i, 0, R) {
    dp[i].resize(C, {0, 0, 0, 0});
    table[i].resize(C, 0);
  }
  REP(i, 0, K) {
    int r, c, v;
    cin >> r >> c >> v;
    r--;
    c--;
    table[r][c] = v;
  }
  dp[0][0][1] = table[0][0];
  REP(j, 1, C) {
    REP(k, 1, 4)
    dp[0][j][k] = max(dp[0][j - 1][k], dp[0][j - 1][k - 1] + table[0][j]);
  }
  REP(i, 1, R) {
    dp[i][0][0] = dp[i - 1][0][1];
    dp[i][0][1] = dp[i - 1][0][1] + table[i][0];
    REP(j, 1, C) {
      dp[i][j][0] = max(dp[i][j - 1][0],
                        *max_element(dp[i - 1][j].begin(), dp[i - 1][j].end()));
      if (table[i][j] <= 0)
        REP(k, 1, 4) dp[i][j][k] = dp[i][j - 1][k];
      else {
        dp[i][j][1] = max(dp[i][j - 1][1],
                          table[i][j] + max(dp[i][j - 1][0],
                                            *max_element(dp[i - 1][j].begin(),
                                                         dp[i - 1][j].end())));
        dp[i][j][2] = max(dp[i][j - 1][2], dp[i][j - 1][1] + table[i][j]);
        dp[i][j][3] = max(dp[i][j - 1][3], dp[i][j - 1][2] + table[i][j]);
      }
    }
  }
  PRI(*max_element(dp[R - 1][C - 1].begin(), dp[R - 1][C - 1].end()));
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <bitset>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdint.h>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;
using ll = long long;
using ull = unsigned long long;
using uint = unsigned int;
using ldouble = long double;

// BDD,ZDD,kdtree, bdtree,bicalc, bolonoy, doloney, tree,
// chinesemod,segmenttree,daikusutora, saidairyuu, 2bugurahu, heirokenshutu,
// topologicalsort, kyourenketuseibun

#define PRI(s) cout << s << endl
#define PRIY PRI("Yes")
#define PRIN PRI("No")

template <typename X> class segtree {
public:
  using func_xx = function<X(X, X)>;
  int N;
  func_xx fxx;
  const X ex;
  vector<X> dat;
  bool changed;

  segtree() = default;
  segtree(const segtree &) = default;
  segtree(int n, func_xx fxx, const X &ex) : fxx(fxx), ex(ex), changed(false) {
    if (n == 1)
      ++n;
    int x = 1;
    while (x < n)
      x *= 2;
    N = x;
    dat.resize(2 * N - 1, ex);
  }

  // 代入後のデータ更新
  void update(int k) {
    dat[k] = fxx(dat[2 * k + 1], dat[2 * k + 2]);
    if (k != 0)
      update((k - 1) / 2);
  }
  void build() {
    for (int i = N - 2; i >= 0; --i)
      dat[i] = fxx(dat[2 * i + 1], dat[2 * i + 2]);
    changed = false;
  }
  // 要素の代入
  void set(int i, const X &x, bool update = true) {
    dat[N - 1 + i] = x;
    if (update)
      this->update((N - 2 + i) / 2);
    else
      changed = true;
  }
  // 要素の参照
  const X &elem(int i) { return dat[N - 1 + i]; }
  //[a,b)間のクエリfxx
  X query_sub(int a, int b, int k, int l, int r) {
    if (r <= a || b <= l) {
      return ex;
    } else if (a <= l && r <= b) {
      return dat[k];
    } else {
      return fxx(query_sub(a, b, 2 * k + 1, l, (l + r) / 2),
                 query_sub(a, b, 2 * k + 2, (l + r) / 2, r));
    }
  }
  X query(int a, int b) {
    if (changed)
      build();
    return query_sub(a, b, 0, 0, N);
  }
};

using arr = array<ll, 3>;
arr e = {0, 0, 0};
auto f = [](arr a, arr b) {
  int ca = 0;
  int cb = 0;
  arr ret = {0, 0, 0};
  for (int i = 0; i < 3; ++i) {
    if (a[ca] > b[cb]) {
      ret[i] = a[ca];
      ++ca;
    } else {
      ret[i] = b[cb];
      ++cb;
    }
  }
  return ret;
};
vector<segtree<arr>> t;
vector<vector<ll>> dp;
int R, C, K;

ll sum(arr a) { return a[0] + a[1] + a[2]; }

int main() {
  cin >> R >> C >> K;
  dp.resize(R);
  for (int i = 0; i < R; ++i) {
    dp[i].resize(C);
    for (int j = 0; j < C; ++j)
      dp[i][j] = 0;
    t.emplace_back(segtree<arr>(C, f, e));
  }
  for (int i = 0; i < K; ++i) {
    ll r, c, v;
    cin >> r >> c >> v;
    --r;
    --c;
    t[r].set(c, arr{v, 0, 0}, false);
  }
  for (int i = 0; i < R; ++i)
    t[i].build();
  for (int j = 0; j < C; ++j) {
    dp[0][j] = sum(t[0].query(0, j + 1));
  }
  for (int i = 1; i < R; ++i) {
    dp[i][0] = dp[i - 1][0] + t[i].elem(0)[0];

    for (int j = 1; j < C; ++j) {
      if (t[i].elem(j)[0] > 0) {
        for (int k = 0; k <= j; ++k)
          dp[i][j] = max(dp[i][j], dp[i - 1][k] + sum(t[i].query(k, j + 1)));
      } else
        dp[i][j] = max(dp[i - 1][j], dp[i][j - 1]);
    }
  }
  PRI(dp[R - 1][C - 1]);
  return 0;
}