1683. Maximum Number of Ones

import numpy as np

height=3
width=3
sidelength=2
maxones=1

key = [[0]*sidelength for _ in range(sidelength)]

def countOnes(rows=sidelength, cols=sidelength):
  ans = 0
  for i in range(rows):
    for j in range(cols):
      ans += key[i][j]
  return ans

if height%sidelength != 0 and width%sidelength != 0:
  count = maxones

  for b in range(sidelength):
    x = 0
    y = b
    while(True):
      if count == 0: break
      if count == 1:
        key[x][y] = 1
        count = 0
        break
      if x==y:
        key[x][y] = 1
        count-=1
        break
      if x-y==-1:
        key[x][y] = key[y][x] = 1
        count-=2
        break

      key[x][y] = key[y][x] = 1
      count -= 2
      x+=1
      y-=1

    if count == 0:
      break

  if count > 0:
    for a in range(1, sidelength):
      x = a
      y = sidelength - 1
      while(True):
        if count == 0: break
        if count == 1:
          key[x][y] = 1
          count = 0
          break
        if x==y:
          key[x][y]=1
          count-=1
          break
        if x-y==-1:
          key[x][y]=key[y][x]=1
          count-=2
          break

        key[x][y] = key[y][x] = 1
        count -= 2
        x+=1
        y-=1

      if count == 0:
        break
else:
  temp = [1] * maxones + [0] * (sidelength*sidelength - maxones)
  if sidelength%width==0:
    for i in range(sidelength):
      for j in range(sidelength):
        key[i][j] = temp[i*sidelength+j]
  else:
     for i in range(sidelength):
       for j in range(sidelength):
         key[j][i] = temp[i*sidelength+j]

res1 = (height//sidelength)*(width//sidelength)*maxones
res2 = (height//sidelength)*countOnes(cols=width%sidelength)
res3 = (width//sidelength)*countOnes(rows=height%sidelength)
res4 = countOnes(rows=height%sidelength, cols=width%sidelength)

res = res1+res2+res3+res4
ans = key

temp=key
for i in range(width//sidelength-1):
  ans = np.append(ans, temp, axis=1)
ans = np.append(ans, np.array(temp)[:,:width%sidelength], axis=1)

temp=ans
for i in range(height//sidelength-1):
  ans = np.append(ans, temp, axis=0)
ans = np.append(ans, np.array(temp)[:height%sidelength], axis=0)

print(ans)
print(res)