public class Solution {
    public static int[] almostMagicSquare(int[] a) {
        int[][] grid = new int[3][3];
        for (int i = 0; i < 9; i++) {
            grid[i / 3][i % 3] = a[i];
        }

        int[] rowSum = new int[3];
        int[] colSum = new int[3];
        int maxSum = 0;

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

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

        for (int i = 0, j = 0; i < 3 && j < 3; ) {
            int 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++;
            }
        }

        int[] result = new int[9];
        for (int i = 0; i < 9; i++) {
            result[i] = grid[i / 3][i % 3];
        }

        return result;
    }
}