Matrix

 1type Matrix struct {
 2	arr      []int
 3	row, col int
 4}
 5
 6func newMatrix(a []int, row, col int) *Matrix {
 7	arr := make([]int, row*col)
 8	copy(arr, a)
 9	return &Matrix{arr, row, col}
10}
11
12func newIMatrix(n int) *Matrix {
13	arr := make([]int, n*n)
14	for i := 0; i < n; i++ {
15		arr[i+i*n] = 1
16	}
17	return &Matrix{arr, n, n}
18}
19
20// At function
21func (m Matrix) At(i, j int) int { return m.arr[i*m.col+j] }
22
23// AddMatrix function
24func AddMatrix(a, b *Matrix) *Matrix {
25	if a.row != b.row || a.col != b.col {
26		return &Matrix{}
27	}
28	res := newMatrix(nil, a.row, a.col)
29	for i := 0; i < a.row*a.col; i++ {
30		res.arr[i] = a.arr[i] + b.arr[i]
31	}
32	return res
33}
34
35// SubMatrix function
36func SubMatrix(a, b *Matrix) *Matrix {
37	if a.row != b.row || a.col != b.col {
38		return &Matrix{}
39	}
40	res := newMatrix(nil, a.row, a.col)
41	for i := 0; i < a.row*a.col; i++ {
42		res.arr[i] = a.arr[i] - b.arr[i]
43	}
44	return res
45}
46
47// MulMatrix function
48func MulMatrix(a, b *Matrix) *Matrix {
49	if a.col != b.row {
50		return &Matrix{}
51	}
52	res := newMatrix(nil, a.row, b.col)
53	for i := 0; i < res.row; i++ {
54		for j := 0; j < res.col; j++ {
55			for k := 0; k < a.col; k++ {
56				res.arr[i*res.col+j] += a.arr[i*a.col+k] * b.arr[k*b.col+j]
57			}
58		}
59	}
60	return res
61}
62
63// PowMatrix function
64func PowMatrix(a *Matrix, n int) *Matrix {
65	if a.row != a.col {
66		return &Matrix{}
67	}
68	t := newMatrix(a.arr, a.row, a.col)
69	res := newIMatrix(a.row)
70	for n > 0 {
71		if n&1 == 1 {
72			res = MulMatrix(res, t)
73		}
74		t = MulMatrix(t, t)
75		n >>= 1
76	}
77	return res
78}