File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 1+ import sys
2+
3+ read = lambda : sys .stdin .readline ().rstrip ()
4+
5+
6+ class Problem :
7+ def __init__ (self ):
8+ self .size , self .exponent = map (int , read ().split ())
9+ self .matrix = [list (map (int , read ().split ())) for _ in range (self .size )]
10+
11+ def solve (self ) -> None :
12+ for row in self .matrix_pow (self .matrix , self .exponent ):
13+ print (* list (map (lambda x : x % 1_000 , row )))
14+
15+ def matrix_pow (self , matrix : list [list [int ]], exponent : int ) -> list [list [int ]]:
16+ if exponent == 1 :
17+ return self .matrix
18+
19+ half = self .matrix_pow (matrix , exponent // 2 )
20+ half = self .mat_mul (half , half )
21+
22+ if exponent % 2 == 1 :
23+ half = self .mat_mul (matrix , half )
24+
25+ return half
26+
27+ def mat_mul (
28+ self ,
29+ matrix_a : list [list [int ]],
30+ matrix_b : list [list [int ]],
31+ ) -> list [list [int ]]:
32+ matrix = [[0 for _ in range (self .size )] for _ in range (self .size )]
33+
34+ for idx in range (self .size ** 2 ):
35+ row , col = idx // self .size , idx % self .size
36+ for index in range (self .size ):
37+ matrix [row ][col ] += matrix_a [row ][index ] * matrix_b [index ][col ]
38+
39+ matrix [row ][col ] %= 1_000
40+
41+ return matrix
42+
43+
44+ if __name__ == "__main__" :
45+ Problem ().solve ()
Original file line number Diff line number Diff line change 1+ [
2+ {
3+ "input" : [
4+ " 2 5" ,
5+ " 1 2" ,
6+ " 3 4"
7+ ],
8+ "expected" : [
9+ " 69 558" ,
10+ " 337 406"
11+ ]
12+ },
13+ {
14+ "input" : [
15+ " 3 3" ,
16+ " 1 2 3" ,
17+ " 4 5 6" ,
18+ " 7 8 9"
19+ ],
20+ "expected" : [
21+ " 468 576 684" ,
22+ " 62 305 548" ,
23+ " 656 34 412"
24+ ]
25+ },
26+ {
27+ "input" : [
28+ " 5 10" ,
29+ " 1 0 0 0 1" ,
30+ " 1 0 0 0 1" ,
31+ " 1 0 0 0 1" ,
32+ " 1 0 0 0 1" ,
33+ " 1 0 0 0 1"
34+ ],
35+ "expected" : [
36+ " 512 0 0 0 512" ,
37+ " 512 0 0 0 512" ,
38+ " 512 0 0 0 512" ,
39+ " 512 0 0 0 512" ,
40+ " 512 0 0 0 512"
41+ ]
42+ },
43+ {
44+ "input" : [
45+ " 2 1" ,
46+ " 1000 1000" ,
47+ " 1000 1000"
48+ ],
49+ "expected" : [
50+ " 0 0" ,
51+ " 0 0"
52+ ]
53+ }
54+ ]
Original file line number Diff line number Diff line change 1+ import json
2+ import os .path
3+ import unittest
4+ from io import StringIO
5+ from unittest .mock import patch
6+
7+ from parameterized import parameterized
8+
9+ from main import Problem
10+
11+
12+ def load_sample (filename : str ):
13+ path = os .path .join (os .path .dirname (os .path .abspath (__file__ )), filename )
14+
15+ with open (path , "r" ) as file :
16+ return [(case ["input" ], case ["expected" ]) for case in json .load (file )]
17+
18+
19+ class TestCase (unittest .TestCase ):
20+ @parameterized .expand (load_sample ("sample.json" ))
21+ def test_case (self , case : str , expected : list [str ]):
22+ # When
23+ with (
24+ patch ("sys.stdin.readline" , side_effect = case ),
25+ patch ("sys.stdout" , new_callable = StringIO ) as output ,
26+ ):
27+ Problem ().solve ()
28+
29+ result = output .getvalue ().rstrip ()
30+
31+ # Then
32+ self .assertEqual ("\n " .join (expected ), result )
33+
34+
35+ if __name__ == "__main__" :
36+ unittest .main ()
You can’t perform that action at this time.
0 commit comments