字节跳动第三题2048AC代码 python

direction= int(input())
A = []
for i in range(4):
    A.append([int(j) for j in input().split()])

def find(A,direction):
    # 1上，2下，3左，4右
    if direction == 1:
        # 往上方滑动，从顶部开始向下合并
        for j in range(4):
            # j表示列，对每一列进行相同的操作
            i = 0
            while i<=2:
                if A[i][j]==A[i+1][j]:
                    A[i][j]*=2
                    A[i+1][j]=0
                    i+=2
                else:
                    i+=1
            # 此时已经和并好了，再将中间的0除去
            # T用来记录每一列中的非0值
            T=[]
            for i in range(4):
                if A[i][j]!=0:
                    T.append(A[i][j])
                    A[i][j]=0
            for k in range(len(T)):
                A[k][j]=T[k]
    if direction == 2:
        # 往下方滑动，从底部向上合并
        for j in range(4):
            # j表示列，对每一列进行相同的操作
            i=3
            while i>=1:
                if A[i][j]==A[i-1][j]:
                    A[i][j]*=2
                    A[i-1][j]=0
                    i-=2
                else:
                    i-=1
            T = []
            # 从下往上记录A的第j列的值
            for i in [3,2,1,0]:
                if A[i][j]!=0:
                    T.append(A[i][j])
                    A[i][j]=0
            # 赋值时，从下往上赋值
            for k in range(len(T)):
                A[3-k][j]=T[k]
    if direction == 3:
        # 往左边滑动
        for i in range(4):
            # i表示列，对每一列进行相同的操作
            j = 0
            while j<=2:
                if A[i][j]==A[i][j+1]:
                    A[i][j]*=2
                    A[i][j+1]=0
                    j+=2
                else:
                    j+=1
            # 此时已经和并好了，再将中间的0除去
            # T用来记录每一列中的非0值
            T=[]
            for j in range(4):
                if A[i][j]!=0:
                    T.append(A[i][j])
                    A[i][j]=0
            for k in range(len(T)):
                A[i][k]=T[k]
    if direction == 4:
        # 往右方滑动，从右部向左合并
        for i in range(4):
            j=3
            while j>=1:
                if A[i][j]==A[i][j-1]:
                    A[i][j]*=2
                    A[i][j-1]=0
                    j-=2
                else:
                    j-=1
            T = []
            for j in [3,2,1,0]:
                if A[i][j]!=0:
                    T.append(A[i][j])
                    A[i][j]=0
            for k in range(len(T)):
                A[i][3-k]=T[k]
    return A
B=find(A,direction)
print(B)