AcWing 3534. 矩阵幂

Signed-off-by: Tategoto Azarasi <2724167997@qq.com>

acwing408.cpp

@@ -1500,4 +1500,64 @@ namespace acwing {
 			return 0;
 		}
 	}// namespace acwing3527
+
+	/**
+	 * 3534. 矩阵幂
+	 */
+	namespace acwing3534 {
+		int main(istream &cin, ostream &cout) {
+			int n, p;
+			cin >> n >> p;
+			/*Matrix mat = Matrix(n);
+			for(int i = 0; i < n; i++) {
+				for(int j = 0; j < n; j++) {
+					cin >> mat[i][j];
+				}
+			}
+			for(int i = 1; i < p; i++) {
+				mat = mat * mat;
+			}
+			for(int i = 0; i < n; i++) {
+				for(int j = 0; j < n; j++) {
+					cout << mat[i][j] << ' ';
+				}
+				cout << endl;
+			}*/
+			vector<Matrix *> mats = vector<Matrix *>(p + 1, nullptr);
+			mats[1]               = new Matrix(n);
+			for(int i = 0; i < n; i++) {
+				for(int j = 0; j < n; j++) {
+					cin >> (*mats[1])[i][j];
+				}
+			}
+			Matrix mat = getMat(mats, p);
+			for(int i = 0; i < n; i++) {
+				for(int j = 0; j < n; j++) {
+					cout << mat[i][j] << ' ';
+				}
+				cout << endl;
+			}
+			return 0;
+		}
+
+		Matrix getMat(vector<Matrix *> &mat, int p) {
+			if(mat[p] != nullptr) {
+				return *mat[p];
+			}
+			Matrix res = Matrix::identity((*mat[1])[1].size());
+			for(int i = 0; p >> i != 0; i++) {
+				int masked = p & (1 << i);
+				if(masked != 0) {
+					if(masked != p) {
+						res = res * getMat(mat, p & (1 << i));
+					} else {
+						Matrix m2 = getMat(mat, masked >> 1);
+						res       = res * m2 * m2;
+					}
+				}
+			}
+			mat[p] = new Matrix(res);
+			return res;
+		}
+	}// namespace acwing3534
 }// namespace acwing

acwing408.h

@@ -5,6 +5,7 @@
 #ifndef PROBLEMSCPP_ACWING408_H
 #define PROBLEMSCPP_ACWING408_H
 
+#include "templates.h"
 #include <iostream>
 #include <map>
 #include <set>
@@ -308,6 +309,14 @@ namespace acwing {
 	namespace acwing3527 {
 		int main(istream & /*cin*/, ostream & /*cout*/);
 	}
+
+	/**
+	 * 3534. 矩阵幂
+	 */
+	namespace acwing3534 {
+		Matrix getMat(vector<Matrix *> &mat, int p);
+		int main(istream & /*cin*/, ostream & /*cout*/);
+	}// namespace acwing3534
 }// namespace acwing
 
 #endif//PROBLEMSCPP_ACWING408_H

acwing408_test.cpp

@@ -1787,4 +1787,21 @@ namespace acwing {
 			          ans);
 		}
 	}// namespace acwing3527
+
+	/**
+ * 3534. 矩阵幂
+ */
+	namespace acwing3534 {
+		TEST(acwing3534, case1) {
+			istringstream in("2 2\n"
+			                 "9 8\n"
+			                 "9 3");
+			auto out = ostringstream();
+			main(in, out);
+			const auto ans = out.str();
+			ASSERT_EQ("153 96 \n"
+			          "108 81 \n",
+			          ans);
+		}
+	}// namespace acwing3534
 }// namespace acwing

templates.cpp

@@ -518,3 +518,42 @@ unsigned UnionFind::count() {
 	}
 	return us.size();
 }
+
+Matrix::Matrix(int n) { mat = vector<vector<int>>(n, vector<int>(n, 0)); }
+
+Matrix::Matrix(const Matrix &m) {
+	mat = vector<vector<int>>(m.mat.size(), vector<int>(m.mat.size(), 0));
+	for(int i = 0; i < m.mat.size(); i++) {
+		for(int j = 0; j < m.mat.size(); j++) {
+			mat[i][j] = m.mat[i][j];
+		}
+	}
+}
+
+Matrix Matrix::operator*(const Matrix &m) const {
+	Matrix ret(mat.size());
+	for(int i = 0; i < mat.size(); i++) {
+		for(int j = 0; j < mat.size(); j++) {
+			for(int k = 0; k < mat.size(); k++) {
+				ret.mat[i][j] += mat[i][k] * m.mat[k][j];
+			}
+		}
+	}
+	return ret;
+}
+
+vector<int> &Matrix::operator[](int i) {
+	return mat[i];
+}
+
+const vector<int> &Matrix::operator[](int i) const {
+	return mat[i];
+}
+
+Matrix Matrix::identity(int n) {
+	Matrix ret(n);
+	for(int i = 0; i < n; i++) {
+		ret.mat[i][i] = 1;
+	}
+	return ret;
+}

templates.h

@@ -101,3 +101,17 @@ class UnionFind {
 	int get_size(int x);
 	unsigned count();
 };
+
+/// \brief 矩阵
+class Matrix {
+private:
+	vector<vector<int>> mat;
+
+public:
+	Matrix(int n);
+	Matrix(const Matrix &m);
+	Matrix operator*(const Matrix &m) const;
+	vector<int> &operator[](int i);
+	const vector<int> &operator[](int i) const;
+	static Matrix identity(int n);
+};