像试试避开循环,用numpy做,还专门考虑了一些比较不规范的情况,但有四道题目一直WA,请教各位大佬可能会是什么问题
import numpy as np
code_input = input().strip().split(' ')
N = int(code_input.pop(0))
code_input = np.array(code_input, dtype=int)
code_input_cumsum = code_input.cumsum()
zero_end_index, zero_end_index_count = np.unique(code_input_cumsum[:-1:2], return_counts=True)
zero_beg_index, zero_beg_index_count = np.unique(code_input_cumsum[1:-1:2], return_counts=True)
code_output = np.zeros(N * N, dtype=int)
code_output[zero_end_index] = zero_end_index_count
code_output[zero_beg_index] = code_output[zero_beg_index] - zero_beg_index_count
print(str(code_output.cumsum().reshape(N, N)).translate({ord(ch):'' for ch in '[] '}))