結果

問題 No.840 ほむほむほむら
コンテスト
ユーザー keymoon
提出日時 2022-04-09 00:23:00
言語 C++17(gcc12)
(gcc 12.3.0 + boost 1.87.0)
結果
RE  
実行時間 -
コード長 5,234 bytes
コンパイル時間 2,950 ms
コンパイル使用メモリ 119,048 KB
最終ジャッジ日時 2025-01-28 16:49:52
ジャッジサーバーID
（参考情報） 		judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample RE * 3
other RE * 25
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <iostream>
#include <algorithm>
#include <vector>
#include <array>
#include <set>

using namespace std;

template<typename T>
struct Component {
  int i;
  int j;
  T elem;
  Component() {}
  Component(int i, int j, T elem) :i(i), j(j), elem(elem) {}
};

template<size_t mat_size, uint32_t mod> // mod は 2^30-1 以下
struct ModMatrix {
  using SparseMat = vector<Component<uint64_t>>;
  using Vector = array<uint64_t, mat_size>;

  array<uint64_t, mat_size* mat_size> d;

  uint64_t& at(const size_t i, const size_t j) { return d[i * mat_size + j]; }
  const uint64_t& at(const size_t i, const size_t j) const { return d[i * mat_size + j]; }

  void sparse_pow(long long n, ModMatrix& res, SparseMat& sparse_notation) {
    if (n == 0) {
      identity(res);
      return;
    }
    ModMatrix tmp;
    sparse_pow(n / 2, tmp, sparse_notation);
    if (n & 1) {
      ModMatrix tmp2;
      mul(tmp, tmp, tmp2);
      mul(tmp2, sparse_notation, res);
    }
    else {
      mul(tmp, tmp, res);
    }
  }
  void sparse_pow(long long n, ModMatrix& res) {
    SparseMat sparse;
    as_sparse(sparse);
    sparse_pow(n, res, sparse);
  }
  ModMatrix sparse_pow(long long n) {
    ModMatrix res;
    sparse_pow(n, res);
    return res;
  }

  void as_sparse(SparseMat& res) {
    res.clear();
    for (int i = 0; i < mat_size; i++) {
      for (int j = 0; j < mat_size; j++) {
        if (at(i, j)) res.emplace_back(i, j, at(i, j));
      }
    }
  }
  SparseMat as_sparse() {
    SparseMat res;
    as_sparse(res);
    return res;
  }


  // 参考: https://ameblo.jp/morimoridiary/entry-10630261646.html
  static void mul(const ModMatrix& A, const ModMatrix& B, ModMatrix& res) {
    const int BLOCK = 16;
    uint64_t temp = 0, * ptemp = 0, partSum[8] = {};
    res.d.fill(0);
    const uint64_t* ptempc = 0;
    for (int i0 = 0; i0 < mat_size; i0 += BLOCK) {
      for (int j0 = 0; j0 < mat_size; j0 += BLOCK) {
        for (int k0 = 0; k0 < mat_size; k0 += BLOCK) {
          for (int i = i0; i < i0 + BLOCK; i++) {
            for (int j = j0; j < j0 + BLOCK; j += 8) {
              partSum[0] = 0;
              partSum[1] = 0;
              partSum[2] = 0;
              partSum[3] = 0;
              partSum[4] = 0;
              partSum[5] = 0;
              partSum[6] = 0;
              partSum[7] = 0;
              ptempc = &B.d[k0 * mat_size + j];
              for (int k = k0; k < k0 + BLOCK; k++) {
                temp = A.d[i * mat_size + k];
                partSum[0] += temp * ptempc[0];
                partSum[1] += temp * ptempc[1];
                partSum[2] += temp * ptempc[2];
                partSum[3] += temp * ptempc[3];
                partSum[4] += temp * ptempc[4];
                partSum[5] += temp * ptempc[5];
                partSum[6] += temp * ptempc[6];
                partSum[7] += temp * ptempc[7];
                ptempc = &ptempc[mat_size];
              }
              ptemp = &res.d[i * mat_size + j];
              (ptemp[0] += partSum[0]) %= mod;
              (ptemp[1] += partSum[1]) %= mod;
              (ptemp[2] += partSum[2]) %= mod;
              (ptemp[3] += partSum[3]) %= mod;
              (ptemp[4] += partSum[4]) %= mod;
              (ptemp[5] += partSum[5]) %= mod;
              (ptemp[6] += partSum[6]) %= mod;
              (ptemp[7] += partSum[7]) %= mod;
            }
          }
        }
      }
    }
  }
  static ModMatrix mul(const ModMatrix& A, const ModMatrix& B) {
    ModMatrix res;
    mul(A, B, res);
    return res;
  }

  static void mul(const ModMatrix& A, const SparseMat& B, ModMatrix& res) {
    res.d.fill(0);
    for (auto&& [j, k, elem] : B) {
      for (int i = 0; i < mat_size; i++) {
        res.at(i, k) += A.at(i, j) * elem;
        res.at(i, k) %= mod;
      }
    }
  }
  static ModMatrix mul(const ModMatrix& A, const SparseMat& B) {
    ModMatrix res;
    mul(A, B, res);
    return res;
  }

  static void mul(const ModMatrix& A, const Vector& B, Vector& res) {
    for (int i = 0; i < mat_size; i++) {
      for (int j = 0; j < mat_size; j++) {
        auto mul = A.at(i, j) * B[j];
        res[i] += mul;
      }
    }
  }
  static Vector mul(const ModMatrix& A, const Vector& B) {
    Vector res;
    mul(A, B, res);
    return res;
  }

  static void identity(ModMatrix& id) {
    for (int i = 0; i < mat_size; i++) id.at(i, i) = 1;
  }
  static ModMatrix identity() {
    ModMatrix res;
    identity(res);
    return res;
  }
};

const int max_k = 5;
const int mat_size = max_k * max_k * max_k;

const int MOD = 998244353;
using Mat = ModMatrix<mat_size, MOD>;

int main() {
  int64_t n, k;
  cin >> n >> k;

  Mat mat{};
  for (int i = 0; i < k; i++) {
    for (int j = 0; j < k; j++) {
      for (int l = 0; l < k; l++) {
        mat.at((i + 1) % k * k * k + j * k + l, i * k * k + j * k + l)++;
        mat.at(i * k * k + (j + i) % k * k + l, i * k * k + j * k + l)++;
        mat.at(i * k * k + j * k + (l + j) % k, i * k * k + j * k + l)++;
      }
    }
  }
  Mat::Vector iv;
  iv[0] = 1;

  auto res_vec = Mat::mul(mat.sparse_pow(n), iv);
  int64_t res = 0;
  for (int i = 0; i < k; i++) {
    for (int j = 0; j < k; j++) {
      res += res_vec.at(i * k * k + j * k);
    }
  }

  cout << res % MOD << endl;
}
0