软件测试面试代码题---计算多维数组的对角线元素和

比如三维数组[[1, 2, 3], [4, 5, 6], [7, 8, 9]]，计算主对角线元素和应该为15。
第一个需要手动构造输入的维数和数据。

def temp(n,data_list):
    sum = 0
    if n <= 0:
        return None
    elif n == 1:
        return data_list[0]
    else:
        for i in range(n):
            for j in range(n):
                if i == j:
                    sum += data_list