Gauss Elimination - linear equations systems solver

gauss.py

# passes https://stepik.org/lesson/9582/step/8?unit=1810

(n,m) = tuple([int(x) for x in input().split()])

e = list()
for i in range(n):
    e.append([float(x) for x in input().split()])

order = list()

for i in range(min(n, m)):
    
    row = -1
    for j in range(0, n):
        if j in order:
            continue
        if abs(e[j][i]) > 0.00000001:
            row = j
            break
        
    if row == -1:
        continue
        
    order.append(row)
    
    for j in range(0, n):
        
        if j in order:
            continue

        t = e[j][i] / e[row][i]
        for k in range(m + 1):
            e[j][k] -= e[row][k] * t

for o in reversed(range(len(order))):
    i = o
    x = order[o]
    
    if abs(e[x][i]) < 0.00000001:
        continue
    
    for j in range(0, n):
        if x == j:
            continue
        
        t = e[j][i] / e[x][i]
        for k in range(m + 1):
            e[j][k] -= e[x][k] * t

ans = list()

for i in range(n):
    c = 0
    for k in range(m):
        if abs(e[i][k]) > 0.00000001:
            c += 1
    if c == 0 and abs(e[i][m]) > 0.00000001:
        print("NO")
        exit()
        
for o in range(len(order)):
    i = o
    x = order[o]
        
    if abs(e[x][i]) < 0.000000001:
        continue
    ans.append(e[x][m] / e[x][i])

if len(ans) == m:
    print("YES")
    for i in ans:
        print(i, end=" ")
else:
    print("INF")