csp认证202212-3 第三题JPEG解码

from math import *
global cnt
def hj(x):
    if x>=255:
        return 255
    elif x<=0:
        return 0
    elif x%1<0.5:
        return x//1
    else:
        return x//1 +1
def arf(x):
    if x==0:
        return sqrt(0.5)
    else:
        return 1
def f1(sy1,sx1):#向左下
    global cnt
    while sy1>=0:
        Q[sx1][sy1]=sm[cnt]
        sx1+=1  #sx是第几个数组，即第几行
        sy1-=1
        cnt+=1
def f2(sy1,sx1):#向右上
    global cnt
    while sx1>=0:
        Q[sx1][sy1]=sm[cnt]
        sx1-=1
        sy1+=1
        cnt+=1

def f3(sy1,sx1):#向左下
    global cnt
    while sx1<=7:
        Q[sx1][sy1]=sm[cnt]
        sx1+=1  #sx是第几个数组，即第几行
        sy1-=1
        cnt+=1
def f4(sy1,sx1):#向右上
    global cnt
    while sy1<=7:
        Q[sx1][sy1]=sm[cnt]
        sx1-=1
        sy1+=1
        cnt+=1

lh=[] #量化
cnt=0
for i in range(8):
    lh.append(list(map(int,input().split())))
n=int(input())
T=int(input())
sm=list(map(int,input().split()))  #扫描数据
for i in range(64-n):
    sm.append(0) #补0
Q=[[0 for i in range(8)] for j in range(8)]

Q[0][0]=sm[cnt]
cnt+=1
x=1
y=0
while x!=7 and y!=8:
    f1(x,y)
    x,y=y,x
    y+=1
    f2(x,y)
    x,y=y,x
    x+=1
f1(x,y)
x,y=y,x
x+=1
#开始向右上走


while cnt<=63:
    f4(x,y)
    x,y=y,x
    y+=1
    f3(x,y)
    x,y=y,x
    x+=1
for i in range(8):
    for j in range(8):
        lh[i][j]*=Q[i][j]

ans=[[0 for i in range(8)] for j in range(8)]
for i in range(8):
    for j in range(8):
        for v in range(8):
            for u in range(8):
                ans[i][j]+=arf(u)*arf(v)*lh[u][v]*cos((pi/8)*(i+0.5)*u)*cos((pi/8)*(j+0.5)*v)

for i in range(8):
    for j in range(8):
        ans[i][j]*=0.25
        ans[i][j]+=128
        ans[i][j]=int(hj(ans[i][j]))
if T==0:
    for i in Q:
        print(*i)
elif T==1:
    for i in lh:
        print(*i)
elif T==2:
    for i in ans:
        print(*i)