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

#include <bits/stdc++.h>

#define REP(i, n) for (int i = 0; i < (int)(n); i++)
#define FOR(i, k, n) for (int i = (k); i < (int)(n); i++)
#define getb(i, j) (((i) >> (j)) & 1)

using namespace std;

const int mod = 1000000007;

struct Mod {
  int num;
  Mod() : num(0) { ; }
  Mod(int n) : num(n) { ; }
  operator int() { return num; }
};

Mod operator+(Mod a, Mod b) { return Mod((a.num + b.num) % mod); }
Mod operator-(Mod a, Mod b) { return Mod((mod + a.num - b.num) % mod); }
Mod operator*(Mod a, Mod b) { return Mod(((long long)a.num * b.num) % mod); }
Mod operator+=(Mod &a, Mod b) { return a = a + b; }
Mod operator-=(Mod &a, Mod b) { return a = a - b; }

typedef array<Mod, 9> Array;
typedef array<Array, 9> Matrix;
typedef array<Matrix, 2> Cube;
typedef array<Cube, 2> DPTable;
typedef tuple<int, int, int, int> J;
int K;

struct UnionFind {
  vector<int> parent;
  UnionFind(int n) { parent.assign(n, -1); }
  int root(int x) { return parent[x] < 0 ? x : parent[x] = root(parent[x]); }
  bool merge(int x, int y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (parent[y] < parent[x])
      swap(x, y);
    if (parent[x] == parent[y])
      --parent[x];
    parent[y] = x;
    return true;
  }
};

DPTable step_normal(DPTable &prev_dp, bool is_cn) {
  DPTable dp;
  REP(i, 2) REP(j, 2) REP(k, 2) REP(l, 2) {
    REP(p1, K + 1) REP(p2, K + 1) {
      if (i == j) {
        if (i == k && j == l) {
          if ((p1 + p2) < K - 1)
            dp[i][j][p1 + p2 + 2][0] += prev_dp[k][l][p1][p2];
        } else if (i == k) {
          if (p1 < K - 1)
            dp[i][j][p1 + 2][0] += prev_dp[k][l][p1][p2];
        } else if (j == l) {
          if (p2 < K - 1)
            dp[i][j][p2 + 2][0] += prev_dp[k][l][p1][p2];
        } else {
          dp[i][j][2][0] += prev_dp[k][l][p1][p2];
        }
      } else {
        if (i == k && j == l) {
          if (p1 < K && p2 < K)
            dp[i][j][p1 + 1][p2 + 1] += prev_dp[k][l][p1][p2];
        } else if (i == k) {
          if (k == l && is_cn) {
            if ((p1 + p2) < K)
              dp[i][j][p1 + p2 + 1][1] += prev_dp[k][l][p1][p2];
          } else {
            if (p1 < K)
              dp[i][j][p1 + 1][1] += prev_dp[k][l][p1][p2];
          }
        } else if (j == l) {
          if (k == l && is_cn) {
            if ((p1 + p2) < K)
              dp[i][j][1][p1 + p2 + 1] += prev_dp[k][l][p1][p2];
          } else {
            if (p2 < K)
              dp[i][j][1][p2 + 1] += prev_dp[k][l][p1][p2];
          }
        } else {
          dp[i][j][1][1] += prev_dp[k][l][p1][p2];
        }
      }
    }
  }
  return dp;
}

struct Node {
  enum class Type { Normal, Edge, Junction } type;
  DPTable dp; // l,r,ln,rn
  vector<shared_ptr<Node>> children;
  bool is_cn() { return type != Type::Edge; }
  Node(const vector<int> &len, const vector<J> &jct, int index, int p, int d) {
    REP(i, 2) REP(j, 2) REP(k, K + 1) REP(l, K + 1) dp[i][j][k][l] = 0;
    if (p == 0 || p == len[index] - 1) { // Edge
      type = Type::Edge;
      if (p + d < 0 || p + d >= len[index]) {
        REP(i, 2) REP(j, 2) dp[i][j][1][1] = 1;
      } else {
        children.push_back(make_shared<Node>(len, jct, index, p + d, d));
        // ||
        // kl
        // ij
        REP(i, 2) REP(j, 2) REP(k, 2) REP(l, 2) {
          UnionFind uf(4);
          if (i == k)
            uf.merge(0, 2);
          if (j == l)
            uf.merge(1, 3);
          if (k == l)
            uf.merge(2, 3);
          REP(m, K + 1) REP(n, K + 1) {
            int w[4] = {1, 1, m, n};
            map<int, int> mp;
            REP(o, 4) mp[uf.root(o)] += w[o];
            int r0 = mp[uf.root(0)];
            int r1 = mp[uf.root(1)];
            if (r0 > K || r1 > K)
              continue;
            dp[i][j][r0][r1] += children.front()->dp[k][l][m][n];
          }
        }
      }
    } else { // Normal, Junction
      bool isj = false;
      REP(i, jct.size()) {
        int f, p1, t, p2;
        tie(f, p1, t, p2) = jct[i];
        if (index == t) {
          swap(f, t);
          swap(p1, p2);
        }
        if (index == f) {
          if (min(p, p + d) == p1) {
            children.push_back(
                make_shared<Node>(len, jct, index, p + 2 * d, d));
            if (d > 0) {
              children.push_back(make_shared<Node>(len, jct, t, p2, -1));
              children.push_back(make_shared<Node>(len, jct, t, p2 + 1, 1));
            } else {
              children.push_back(make_shared<Node>(len, jct, t, p2 + 1, 1));
              children.push_back(make_shared<Node>(len, jct, t, p2, -1));
            }
            isj = true;
          }
        }
      }
      if (isj) { // Junction
        // ||
        // 45
        //-12-
        //-03-
        // ||
        REP(i, 1 << 5) {
          const int bits = i << 1;
          const bool iscnc = children.front()->is_cn();
          bitset<6> bs(bits);
          UnionFind uf(6);
          int ary[6][2] = {{0, 1}, {2, 3}, {0, 3}, {1, 2}, {1, 4}, {2, 5}};
          REP(j, 6)
          if (bs[ary[j][0]] == bs[ary[j][1]])
            uf.merge(ary[j][0], ary[j][1]);
          if (iscnc && bs[4] == bs[5])
            uf.merge(4, 5);
          int ufr[6];
          REP(j, 6) ufr[j] = uf.root(j);
          int powK1[7];
          powK1[0] = 1;
          REP(j, 6) powK1[j + 1] = powK1[j] * (K + 1);
          int loop = powK1[6];
          if (iscnc && bs[4] == bs[5])
            loop = powK1[5];
          REP(j, loop) {
            int pi[6];
            int cnt[6] = {0, 0, 0, 0, 0, 0};
            int ind = j;
            REP(k, 6) {
              pi[k] = ind % (K + 1);
              ind /= K + 1;
              cnt[ufr[k]] += pi[k];
            }
            if (!getb(bits, 1) && pi[0]) {
              j = (j / powK1[1]) * powK1[1] + (powK1[1] - 1);
              continue;
            } else if (getb(bits, 2) == getb(bits, 3) && pi[2]) {
              j = (j / powK1[3]) * powK1[3] + (powK1[3] - 1);
              continue;
            } else if (!iscnc) {
              if (pi[4] > 1) {
                j = (j / powK1[5]) * powK1[5] + (powK1[5] - 1);
                continue;
              } else if (pi[5] > 1) {
                break;
              }
            }
            bool ok = true;
            REP(k, 6) if (cnt[k] > K) ok = false;
            if (!ok) {
              int k = 1;
              while ((j % powK1[k]) == 0)
                ++k;
              j = (j / powK1[k]) * powK1[k] + (powK1[k] - 1);
              continue;
            }
            Mod addv =
                children[0]->dp[getb(bits, 4)][getb(bits, 5)][pi[4]][pi[5]] *
                children[1]->dp[getb(bits, 1)][0][pi[1]][pi[0]] *
                children[2]->dp[getb(bits, 3)][getb(bits, 2)][pi[3]][pi[2]];
            if (!getb(bits, 3))
              dp[0][getb(bits, 3)][cnt[ufr[0]]][0] += addv;
            else
              dp[0][getb(bits, 3)][cnt[ufr[0]]][cnt[ufr[3]]] += addv;
          }
        }
        REP(i, 2) REP(j, K + 1) REP(k, K + 1) {
          dp[1][i][j][k] = dp[0][1 - i][j][k];
        }
      } else { // Normal
        children.push_back(make_shared<Node>(len, jct, index, p + d, d));
        auto child = children.front();
        dp = step_normal(child->dp, child->is_cn());
      }
    }
  }
  void out(int dep = 1) {
    cout << string(dep, ' ') << 'o' << endl;
    for (auto child : children)
      child->out(dep + 1);
  }
};

int main() {
  while (1) {
    int n;
    cin >> n >> K;
    if (!n)
      break;
    vector<int> l(n);
    REP(i, n) {
      cin >> l[i];
      l[i] += 3;
    }
    vector<J> jct;
    REP(i, n - 1) {
      int f, p, t, q;
      cin >> f >> p >> t >> q;
      ++p;
      ++q;
      jct.emplace_back(f, p, t, q);
    }
    auto root = Node(l, jct, 0, 0, 1);
    Mod sum = 0;
    REP(i, 2) REP(j, 2) REP(k, K + 1) REP(l, K + 1) {
      sum += root.dp[i][j][k][l];
    }
    cout << sum << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

#define REP(i, n) for (int i = 0; i < (int)(n); i++)
#define FOR(i, k, n) for (int i = (k); i < (int)(n); i++)

using namespace std;

const int mod = 1000000007;

struct Mod {
  int num;
  Mod() : num(0) { ; }
  Mod(int n) : num(n) { ; }
  operator int() { return num; }
};

Mod operator+(Mod a, Mod b) { return Mod((a.num + b.num) % mod); }
Mod operator-(Mod a, Mod b) { return Mod((mod + a.num - b.num) % mod); }
Mod operator*(Mod a, Mod b) { return Mod(((long long)a.num * b.num) % mod); }
Mod operator+=(Mod &a, Mod b) { return a = a + b; }
Mod operator-=(Mod &a, Mod b) { return a = a - b; }

typedef array<Mod, 9> Array;
typedef array<Array, 9> Matrix;
typedef array<Matrix, 2> Cube;
typedef array<Cube, 2> DPTable;
typedef tuple<int, int, int, int> J;
int K;

struct UnionFind {
  vector<int> parent;
  UnionFind(int n) { parent.assign(n, -1); }
  int root(int x) { return parent[x] < 0 ? x : parent[x] = root(parent[x]); }
  bool merge(int x, int y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (parent[y] < parent[x])
      swap(x, y);
    if (parent[x] == parent[y])
      --parent[x];
    parent[y] = x;
    return true;
  }
};

DPTable step_normal(DPTable &prev_dp, bool is_cn) {
  DPTable dp;
  REP(i, 2) REP(j, 2) REP(k, 2) REP(l, 2) {
    REP(p1, K + 1) REP(p2, K + 1) {
      if (i == j) {
        if (i == k && j == l) {
          if ((p1 + p2) < K - 1)
            dp[i][j][p1 + p2 + 2][0] += prev_dp[k][l][p1][p2];
        } else if (i == k) {
          if (p1 < K - 1)
            dp[i][j][p1 + 2][0] += prev_dp[k][l][p1][p2];
        } else if (j == l) {
          if (p2 < K - 1)
            dp[i][j][p2 + 2][0] += prev_dp[k][l][p1][p2];
        } else {
          dp[i][j][2][0] += prev_dp[k][l][p1][p2];
        }
      } else {
        if (i == k && j == l) {
          if (p1 < K && p2 < K)
            dp[i][j][p1 + 1][p2 + 1] += prev_dp[k][l][p1][p2];
        } else if (i == k) {
          if (k == l && is_cn) {
            if ((p1 + p2) < K)
              dp[i][j][p1 + p2 + 1][1] += prev_dp[k][l][p1][p2];
          } else {
            if (p1 < K)
              dp[i][j][p1 + 1][1] += prev_dp[k][l][p1][p2];
          }
        } else if (j == l) {
          if (k == l && is_cn) {
            if ((p1 + p2) < K)
              dp[i][j][1][p1 + p2 + 1] += prev_dp[k][l][p1][p2];
          } else {
            if (p2 < K)
              dp[i][j][1][p2 + 1] += prev_dp[k][l][p1][p2];
          }
        } else {
          dp[i][j][1][1] += prev_dp[k][l][p1][p2];
        }
      }
    }
  }
  return dp;
}

struct Node {
  enum class Type { Normal, Edge, Junction } type;
  DPTable dp; // l,r,ln,rn
  vector<shared_ptr<Node>> children;
  bool is_cn() { return type != Type::Edge; }
  Node(const vector<int> &len, const vector<J> &jct, int index, int p, int d) {
    REP(i, 2) REP(j, 2) REP(k, K + 1) REP(l, K + 1) dp[i][j][k][l] = 0;
    if (p == 0 || p == len[index] - 1) { // Edge
      type = Type::Edge;
      if (p + d < 0 || p + d >= len[index]) {
        REP(i, 2) REP(j, 2) dp[i][j][1][1] = 1;
      } else {
        children.push_back(make_shared<Node>(len, jct, index, p + d, d));
        // ||
        // kl
        // ij
        REP(i, 2) REP(j, 2) REP(k, 2) REP(l, 2) {
          UnionFind uf(4);
          if (i == k)
            uf.merge(0, 2);
          if (j == l)
            uf.merge(1, 3);
          if (k == l)
            uf.merge(2, 3);
          REP(m, K + 1) REP(n, K + 1) {
            int w[4] = {1, 1, m, n};
            map<int, int> mp;
            REP(o, 4) mp[uf.root(o)] += w[o];
            int r0 = mp[uf.root(0)];
            int r1 = mp[uf.root(1)];
            if (r0 > K || r1 > K)
              continue;
            if (uf.root(0) == uf.root(1)) {
              dp[i][j][r0][0] += children.front()->dp[k][l][m][n];
            } else {
              dp[i][j][r0][r1] += children.front()->dp[k][l][m][n];
            }
          }
        }
      }
    } else { // Normal, Junction
      bool isj = false;
      int n = len.size();
      REP(i, n - 1) {
        int f, p1, t, p2;
        tie(f, p1, t, p2) = jct[i];
        if (index == t) {
          swap(f, t);
          swap(p1, p2);
        }
        if (index == f) {
          if (min(p, p + d) == p1) {
            children.push_back(
                make_shared<Node>(len, jct, index, p + 2 * d, d));
            if (d > 0) {
              children.push_back(make_shared<Node>(len, jct, t, p2, -1));
              children.push_back(make_shared<Node>(len, jct, t, p2 + 1, 1));
            } else {
              children.push_back(make_shared<Node>(len, jct, t, p2 + 1, 1));
              children.push_back(make_shared<Node>(len, jct, t, p2, -1));
            }
            isj = true;
          }
        }
      }
      if (isj) { // Junction
        // cout << "J" << endl;
        //  ||
        //  45
        //-12-
        //-03-
        //  ||
        REP(i, 1 << 5) {
          bitset<6> bs(i << 1);
          UnionFind uf(6);
          int ary[6][2] = {{0, 1}, {2, 3}, {0, 3}, {1, 2}, {1, 4}, {2, 5}};
          REP(j, 6)
          if (bs[ary[j][0]] == bs[ary[j][1]])
            uf.merge(ary[j][0], ary[j][1]);
          if (children.front()->is_cn() && bs[4] == bs[5])
            uf.merge(4, 5);
          int powK1[7];
          powK1[0] = 1;
          REP(j, 6) powK1[j + 1] = powK1[j] * (K + 1);
          int loop = powK1[6];
          if (children[0]->is_cn() && bs[4] == bs[5])
            loop = powK1[5];
          REP(j, loop) {
            int pi[6];
            int cnt[6] = {0, 0, 0, 0, 0, 0};
            int ind = j;
            REP(k, 6) {
              pi[k] = ind % (K + 1);
              ind /= K + 1;
              cnt[uf.root(k)] += pi[k];
            }
            if (!bs[1] && pi[1]) {
              j = (j / powK1[2]) * powK1[2] + (powK1[2] - 1);
              continue;
            } else if (bs[2] == bs[3] && pi[3]) {
              j = (j / powK1[4]) * powK1[4] + (powK1[4] - 1);
              continue;
            }
            if (!children[0]->is_cn()) {
              if (pi[4] > 1) {
                j = (j / powK1[5]) * powK1[5] + (powK1[5] - 1);
                continue;
              } else if (pi[5] > 1) {
                break;
              }
            }
            bool ok = true;
            REP(k, 6) if (cnt[k] > K) ok = false;
            if (!ok) {
              int k = 1;
              while ((j % powK1[k]) == 0)
                ++k;
              j = (j / powK1[k]) * powK1[k] + (powK1[k] - 1);
              continue;
            }
            Mod addv = children[0]->dp[bs[4]][bs[5]][pi[4]][pi[5]] *
                       children[1]->dp[0][bs[1]][pi[0]][pi[1]] *
                       children[2]->dp[bs[2]][bs[3]][pi[2]][pi[3]];
            if (!bs[3])
              dp[0][bs[3]][cnt[uf.root(0)]][0] += addv;
            else
              dp[0][bs[3]][cnt[uf.root(0)]][cnt[uf.root(3)]] += addv;
          }
        }
        REP(i, 2) REP(j, K + 1) REP(k, K + 1) {
          dp[1][i][j][k] = dp[0][1 - i][j][k];
        }
      } else { // Normal
        children.push_back(make_shared<Node>(len, jct, index, p + d, d));
        auto child = children.front();
        dp = step_normal(child->dp, child->is_cn());
      }
    }
  }
  void out(int dep = 1) {
    cout << string(dep, ' ') << 'o' << endl;
    for (auto child : children)
      child->out(dep + 1);
  }
};

int main() {
  while (1) {
    int n;
    cin >> n >> K;
    if (!n)
      break;
    vector<int> l(n);
    REP(i, n) {
      cin >> l[i];
      l[i] += 3;
    }
    vector<J> jct;
    REP(i, n - 1) {
      int f, p, t, q;
      cin >> f >> p >> t >> q;
      ++p;
      ++q;
      jct.emplace_back(f, p, t, q);
    }
    auto root = Node(l, jct, 0, 0, 1);
    // root.out();
    Mod sum = 0;
    REP(i, 2) REP(j, 2) REP(k, K + 1) REP(l, K + 1) {
      sum += root.dp[i][j][k][l];
    }
    cout << sum << endl;
  }
  return 0;
}