matrixkotlinfunctional-programmingmappingcolumn-sum

Kotlin: Elegant way to sum elements across the rows in matrix?


Given an ArrayList (rows) of ArrayList (elements in a row), is there a better way to obtain a list of sums like (row[0][0] + row[1][0] + .... row[n][0]) except writing a "for" loop?


Solution

  • You can use sum and sumBy to achieve this.

    Following are 3 methods: (1) to calculate the overall matrix sum, (2) to calculate the sum for a given row, and (3) to calculate the sum for a given column:

    val arr = arrayListOf(
            arrayListOf(1, 2, 3),
            arrayListOf(4, 5, 6),
            arrayListOf(7, 8, 9)
    )
    
    fun matrixSum() = arr.sumBy { it.sum() }
    fun rowSum(row: Int) = arr[row].sum()
    fun colSum(col: Int) = arr.sumBy {row -> row[col] }
    // shorter alternative: fun colSum(col: Int) = arr.sumBy { it[col] }
    

    According to your question, you would like to call colSum(0). This will extract row[0] (each row's value for index 0) and sum these values.

    EDIT:

    Regarding your question in the comment section, following two methods: (1) To calculate all row sums, returned as a list; and (2) To calculate all col sums, returned as a list.

    fun rowSums() = arr.map { it.sum() }
    fun colSums() = arr[0].indices.map { col -> colSum(col) }
    // shorter alternative:
    // fun colSums() = arr[0].indices.map(::colSum)
    // another version that does not call colSum:
    // fun colSums() = arr[0].indices.map { col -> arr.sumBy { row -> row[col] } }
    

    To explain colSums (which assumes that all rows have the same number of columns): First, you grab your first row (arr[0]), then you get its indices (i.e., its column numbers 0, 1, ..., up to nCols), and then you map these indices to their respective colSum() results. This may not be the most performant way, but it's quite readable.