TreeMapView.kt

package com.pt.treemap

import android.content.Context
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Paint
import android.graphics.Rect
import android.util.AttributeSet
import android.view.View
import com.pt.treemap.TreeMapView.Orientation.HORIZONTAL
import com.pt.treemap.TreeMapView.Orientation.VERTICAL
import kotlin.random.Random

class TreeMapView @JvmOverloads constructor(
    context: Context,
    attrs: AttributeSet? = null,
    defStyleAttr: Int = 0
) : View(context, attrs, defStyleAttr) {

    private var data = emptyList<Int>()
    private var bound: Rect = Rect(0, 0, 0, 0)
    private val random = Random(System.currentTimeMillis())

    private val paint = Paint(Paint.ANTI_ALIAS_FLAG).apply {
        style = Paint.Style.FILL
    }

    private fun randomColor(): Int =
        Color.argb(
            128,
            random.nextInt(0, 255),
            random.nextInt(0, 255),
            random.nextInt(0, 255)
        )

    fun setData(list: List<Int>) {
        data = list
        invalidate()
    }

    override fun onDraw(canvas: Canvas) {
        super.onDraw(canvas)
        if (data.isNotEmpty())
            drawTreeMap(canvas, data)
    }

    override fun onSizeChanged(w: Int, h: Int, oldw: Int, oldh: Int) {
        super.onSizeChanged(w, h, oldw, oldh)
        bound = Rect(0, 0, w, h)
    }

    private fun drawTreeMap(canvas: Canvas, data: List<Int>) {
        require(data.isNotEmpty())
        drawTreeMapRecursive(canvas, data, bound, VERTICAL)
    }

    private fun drawTreeMapRecursive(
        canvas: Canvas,
        data: List<Int>,
        bound: Rect,
        orientation: Orientation
    ) {
        if (data.size == 1) {
            canvas.drawRect(bound, paint.apply {
                color = randomColor()
            })
            canvas.drawText(
                data[0].toString(),
                bound.left.toFloat() + 4,
                bound.bottom.toFloat() - 4,
                paint.apply {
                    color = Color.BLACK
                    textSize = 40f
                })
        } else {
            if (orientation == HORIZONTAL) {
                val size = bound.right - bound.left
                val (left, right) = divList(data)
                val leftSize = size * left.sum() / data.sum()
                val divider = bound.left + leftSize
                val leftBound = Rect(bound.left, bound.top, divider, bound.bottom)
                val rightBound = Rect(divider, bound.top, bound.right, bound.bottom)
                drawTreeMapRecursive(canvas, left, leftBound, VERTICAL)
                drawTreeMapRecursive(canvas, right, rightBound, VERTICAL)
            } else {
                val size = bound.bottom - bound.top
                val (top, bottom) = divList(data)
                val topSize = size * top.sum() / data.sum()
                val divider = bound.top + topSize
                val topBound = Rect(bound.left, bound.top, bound.right, divider)
                val bottomBound = Rect(bound.left, divider, bound.right, bound.bottom)
                drawTreeMapRecursive(canvas, top, topBound, HORIZONTAL)
                drawTreeMapRecursive(canvas, bottom, bottomBound, HORIZONTAL)
            }
        }
    }

    enum class Orientation {
        HORIZONTAL, VERTICAL
    }

    companion object {
        fun divList(data: List<Int>): Pair<List<Int>, List<Int>> {
            require(data.size > 1)
            val total = data.sum()
            val half = total / 2
            var sum = 0
            var i = 0
            while (i < data.size) {
                val value = data[i]
                if (sum + value > half)
                    break
                sum += value
                i += 1
            }
            val rest = total - sum - data[i]
            val position = if (sum + data[i] - rest < rest + data[i] - sum)
                i + 1
            else
                i

            return data.subList(0, position) to data.subList(position, data.size)
        }
    }
}