function almostMagicSquare(a: number[]): number[] {
  const rowSum = [0, 0, 0];
  const colSum = [0, 0, 0];
  let maxSum = 0;

  const grid: number[][] = [a.slice(0, 3), a.slice(3, 6), a.slice(6, 9)];

  for (let i = 0; i < 3; i++) {
    for (let j = 0; j < 3; j++) {
      rowSum[i] += grid[i][j];
      colSum[i] += grid[j][i];
    }
  }

  for (let k = 0; k < 3; k++) {
    maxSum = Math.max(maxSum, rowSum[k]);
    maxSum = Math.max(maxSum, colSum[k]);
  }

  for (let i = 0, j = 0; i < 3 && j < 3; ) {
    const diff = Math.min(maxSum - rowSum[i], maxSum - colSum[j]);
    grid[i][j] += diff;
    rowSum[i] += diff;
    colSum[j] += diff;

    if (rowSum[i] === maxSum) {
      i++;
    }
    if (colSum[j] === maxSum) {
      j++;
    }
  }

  return grid.flat();
}