Показывать изображения в виде коллекции в формате сетки различного размера

Question

Я хочу показывать изображения в виде коллекции в формате сетки различных размеров со знаком плюс (т.е. + 2, + 3 и т. Д.).Это будет выглядеть так ...

Я узнал, что с помощью UICollectionViewFlowLayout этого можно достичь.Поэтому я сослался на этот вопрос переполнения стека и попытался добавить эту строку кода.Но это не сработало

func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, sizeForItemAt indexPath: IndexPath) -> CGSize {

        switch indexPath.item {
        case 0,1:
            return CGSize(width: (UIScreen.main.bounds.width - 16) / 2, height: (UIScreen.main.bounds.width - 16) / 2)
        default:
            return CGSize(width: (UIScreen.main.bounds.width - 32) / 3, height:  (UIScreen.main.bounds.width) / 3)
        }
    }

Кроме того, если количество выбранных изображений равно 2 или 3, то эти изображения должны быть упорядочены следующим образом.

Как сделать так, чтобы изображения выглядели так? ...