Repository: android/compose-samples
File: JetNews/app/src/main/java/com/example/jetnews/ui/interests/InterestsScreen.kt (Kotlin)

Code before:
    onTopicSelect: (String) -> Unit
) {
    BoxWithConstraints {
        val itemMaxWidth = rememberItemMaxWidth(windowMaxWidth = maxWidth, columns = 1)
        val topicModifier = Modifier
            .fillMaxWidth()
            .wrapContentWidth(Alignment.CenterHorizontally)
            .widthIn(max = itemMaxWidth)

        LazyColumn(
            modifier = Modifier
                .padding(top = 16.dp)
                .navigationBarsPadding()
        ) {
            items(topics) { topic ->
                TopicItem(
                    itemTitle = topic,
                    selected = selectedTopics.contains(topic),
                    onToggle = { onTopicSelect(topic) },
                    modifier = topicModifier
                )
                TopicDivider(topicModifier)
            }
        }
    }
}

/**
 * Display a sectioned list of topics
 *
 * @param sections (state) topics to display, grouped by sections
 * @param selectedTopics (state) currently selected topics
 * @param onTopicSelect (event) request a topic+section selection be changed
 */
@Composable
private fun TabWithSections(
    sections: List<InterestSection>,
    selectedTopics: Set<TopicSelection>,
    onTopicSelect: (TopicSelection) -> Unit
) {
    BoxWithConstraints {
        val windowSize = remember(maxWidth) { getWindowSize(maxWidth) }
        val columns = remember(windowSize) { if (windowSize == WindowSize.Compact) 1 else 2 }

        val itemMaxWidth = rememberItemMaxWidth(windowMaxWidth = maxWidth, columns = columns)
        // Group interest sections given the number of columns that fill the screen.
        // Sections are distributed evenly on the screen using the modulus operator.
        val groupedSections: Array<MutableList<InterestSection>> = remember(columns) {
            val sectionsInColumns = Array<MutableList<InterestSection>>(columns) { mutableListOf() }
            sections.forEachIndexed { index, section ->
                val column = index % columns

Review comment (suggestion):
One thing that's a little strange with the current logic: the number of items in each column isn't really balanced. If all of the even sections have 20 items and all of the odd sections have 2 items, then they could be really unbalanced, which is a bit strange.

Code after:
 */
@Composable
private fun TabWithSections(
    sections: List<InterestSection>,
    selectedTopics: Set<TopicSelection>,
    onTopicSelect: (TopicSelection) -> Unit
) {
    BoxWithConstraints {
        val windowSize = remember(maxWidth) { getWindowSize(maxWidth) }
        val columns = remember(windowSize) { if (windowSize == WindowSize.Compact) 1 else 2 }
        val itemMaxWidth = rememberItemMaxWidth(maxWidth, columns)
        val groupedSections = rememberSectionsGroupedInColumns(sections, columns)

        Row(
            horizontalArrangement = Arrangement.SpaceEvenly,
            modifier = Modifier
                .navigationBarsPadding()
                .fillMaxSize()
        ) {
            groupedSections.forEach { sectionsInGroup ->
                LazyColumn(
                    Modifier
                        .widthIn(max = itemMaxWidth)
                        .padding(horizontal = if (columns > 1) 8.dp else 0.dp)
                ) {
                    sectionsInGroup.forEach { (section, topics) ->
                        item {
                            Text(
                                text = section,
                                modifier = Modifier
                                    .padding(16.dp)
                                    .semantics { heading() },
                                style = MaterialTheme.typography.subtitle1
                            )
                        }
                        items(topics) { topic ->
                            TopicItem(
                                itemTitle = topic,
                                selected = selectedTopics.contains(TopicSelection(section, topic)),
                                onToggle = { onTopicSelect(TopicSelection(section, topic)) },
                                modifier = Modifier
                            )
                            TopicDivider(Modifier)
                        }
                    }
                }
            }
        }
    }
}