How to fill a matrix with equal sum of rows and columns? - python-3.x

Question

I have a N x N matrix with integer elements.
We have two inputs : n and k.
There is two condition for solving this problem:
1- sum of matrix's columns and rows should be equal to k.
2- Difference between max and min numbers in matrix should be minimum.
I wrote a code in python but it doesn't work well.
n , k = map(int,input().split())
matrix = [[k//n]*n for i in range(n)]
def row_sum(matrix,row):
return sum(matrix[row])
def col_sum(matrix,col):
res = 0
for i in matrix:
res += i[col]
return res
for i in range(n):
for j in range(n):
if (row_sum(matrix,i) != k) and (col_sum(matrix, j) != k):
matrix[i][j] += 1
for i in matrix:
print(*i)
for example we have a 5x5 matrix that sum of its columns and rows should be equal to 6:
input : 5 6
output :
2 1 1 1 1
1 2 1 1 1
1 1 2 1 1
1 1 1 2 1
1 1 1 1 2
but it doesn't work well:
input : 6 11
output:
2 2 2 2 2 1
2 2 2 2 2 1
2 2 2 2 2 1
2 2 2 2 2 1
2 2 2 2 2 1
1 1 1 1 1 2
I spend a lot of time on this and i can't solve it. Please Help!
(This problem is not a homework or something like that. It's a question from an algorithm contest and the contest is over!)

Answer 1

The solution is to work out the first row (using the code you already have), and then set each row to be the row above it rotated one position.
So for example if the first row has the values
a b c d e
then you rotate one position each row to get
a b c d e
b c d e a
c d e a b
d e a b c
e a b c d
Since each value gets placed in each column once the columns will contain one of each value and so add up to the same total, and since each row has the same values just moved around all the rows add up the same too.
Code:
n , k = map(int,input().split())
matrix = [[k//n]*n for i in range(n)]
def row_sum(matrix,row):
return sum(matrix[row])
for j in range(n):
if (row_sum(matrix,0) != k):
matrix[0][j] += 1
for i in range(1, n):
for j in range(n):
matrix[i][j] = matrix[i-1][(j+1)%n]
for i in matrix:
print(*i)