Jetpack compose Staggered Vertical Grid wrong item selected after reordering

Hi I am using a Jetpack Compose implementation of a Staggered Vertical Grid. It works well but when I reorder items by priority I got wrong selection when I click on some item.

import androidx.compose.runtime.Composable
import androidx.compose.ui.Modifier
import androidx.compose.ui.layout.Layout
import androidx.compose.ui.unit.Dp
import kotlin.math.ceil

@Composable
fun StaggeredVerticalGrid(
    modifier: Modifier = Modifier,
    maxColumnWidth: Dp,
    children: @Composable () -> Unit
) {
    Layout(
        content = children,
        modifier = modifier
    ) { measurables, constraints ->
        check(constraints.hasBoundedWidth) {
            "Unbounded width not supported"
        }
        val columns = ceil(constraints.maxWidth / maxColumnWidth.toPx()).toInt()
        val columnWidth = constraints.maxWidth / columns
        val itemConstraints = constraints.copy(maxWidth = columnWidth)
        val colHeights = IntArray(columns) { 0 } // track each column's height
        val placeables = measurables.map { measurable ->
            val column = shortestColumn(colHeights)
            val placeable = measurable.measure(itemConstraints)
            colHeights[column] += placeable.height
            placeable
        }

        val height = colHeights.maxOrNull()?.coerceIn(constraints.minHeight, constraints.maxHeight)
            ?: constraints.minHeight
        layout(
            width = constraints.maxWidth,
            height = height
        ) {
            val colY = IntArray(columns) { 0 }
            placeables.forEach { placeable ->
                val column = shortestColumn(colY)
                placeable.place(
                    x = columnWidth * column,
                    y = colY[column]
                )
                colY[column] += placeable.height
            }
        }
    }
}

private fun shortestColumn(colHeights: IntArray): Int {
    var minHeight = Int.MAX_VALUE
    var column = 0
    colHeights.forEachIndexed { index, height ->
        if (height < minHeight) {
            minHeight = height
            column = index
        }
    }
    return column
}

I am using this implementation in example below:

@ExperimentalAnimationApi
@Composable
fun ArticleList(
    listOfArticles: List<Article>?,
    onArticleClicked: (Article) -> Unit,
    onRemoveArticle: (Article) -> Unit,
    onDialogShown: () -> Unit
) {
    if (listOfArticles.isNullOrEmpty()) {
        NoArticle(stringResource(R.string.no_article_available))
        return
    }
    StaggeredVerticalGrid(
        maxColumnWidth = 220.dp,
        modifier = Modifier.padding(4.dp)
    ) {
        listOfArticles.forEach { article ->
            AnimatedVisibility(
                visible = true,
                enter = expandVertically(),
                exit = shrinkVertically()
            ) {
                ArticleItem(
                    article = article,
                    onArticleClicked = onArticleClicked,
                    onRemoveArticle = onRemoveArticle,
                    onDialogShown = onDialogShown
                )
            }
        }
    }
}

After searching, I found out I need to add a key as LazyColumn does. How can I achieve this with my implementation? (Sorry for poor English)