Swift中高效提取数组唯一与重复元素的优化方案探讨

需求 #

• 同时获取唯一元素（保留首次出现的实例）与重复元素集合，而非仅去除重复项
• 按元素的id属性对重复元素进行分组

现有实现 #

实现一：单次遍历，空间占用略高 #

该方案通过一次遍历完成所有逻辑，需要维护额外的ID映射字典，但无需二次遍历原数组：

public extension Array where Element: Identifiable {
    /// 过滤数组中的重复元素，返回唯一元素列表和按ID分组的重复元素集合
    /// - Returns: 唯一元素数组、按ID分组的重复元素字典（包含该ID下所有元素）
    func filterDuplicates() -> (uniqueElements: [Element], duplicateElementGroups: [Element.ID: [Element]]) {
        var uniqueElements = [Element]()
        var uniqueIDs = [Element.ID: Element]()
        var duplicateElements = [Element.ID: [Element]]()
        
        for element in self {
            if let existingElement = uniqueIDs[element.id] {
                if var duplicates = duplicateElements[element.id] {
                    duplicates.append(element)
                    duplicateElements[element.id] = duplicates
                } else {
                    duplicateElements[element.id] = [existingElement, element]
                }
            } else {
                uniqueIDs[element.id] = element
                uniqueElements.append(element)
            }
        }
        
        return (uniqueElements, duplicateElements)
    }
}

实现二：两次遍历，可读性更强、空间更优 #

该方案先通过Dictionary(grouping:by:)完成分组，再遍历原数组收集唯一元素，代码逻辑更直观：

public extension Array where Element: Identifiable {
    func filterDuplicates() -> (uniqueElements: [Element], duplicateElementGroups: [Element.ID: [Element]]) {
        var uniqueElements = [Element]()
        var duplicateElements = [Element.ID: [Element]]()
        
        let groupedById = Dictionary(grouping: self, by: \.id)
        
        for element in self {
            if duplicateElements[element.id] == nil {
                uniqueElements.append(element)
            }
            
            if let group = groupedById[element.id], group.count > 1 {
                duplicateElements[element.id] = group
            }
        }
        
        return (uniqueElements, duplicateElements)
    }
}

测试用例 #

通过以下测试验证功能正确性：

// 测试用的Identifiable结构体
struct SampleIdentifiableElement: Identifiable {
    let id: String
    let title: String
}

func test_whenFilterDuplicatesInArray_uniqueAndDuplicateItemsReturns() {
    // 测试数据
    let sampleElements: [SampleIdentifiableElement] = [
        .init(id: "1", title: "title 1"),
        .init(id: "2", title: "title 2 - first duplicate"),
        .init(id: "2", title: "title 2 - second duplicate"),
        .init(id: "2", title: "title 2 - third duplicate"),
        .init(id: "3", title: "title 3"),
        .init(id: "4", title: "title 4 - first duplicate"),
        .init(id: "4", title: "title 4 - second duplicate")
    ]

    // 执行方法
    let (uniqueElements, duplicateElementGroups) = sampleElements.filterDuplicates()
            
    // 验证结果
    XCTAssertEqual(uniqueElements.count, 4)
    XCTAssertEqual(duplicateElementGroups.count, 2)
    XCTAssertEqual(duplicateElementGroups["2"]?.count, 3)
    XCTAssertEqual(duplicateElementGroups["4"]?.count, 2)
    
    XCTAssertEqual(duplicateElementGroups["2"]?.first?.title, "title 2 - first duplicate")
    XCTAssertEqual(duplicateElementGroups["2"]?.last?.title, "title 2 - third duplicate")
            
    XCTAssertEqual(uniqueElements.first?.title, "title 1")
    XCTAssertEqual(uniqueElements.last?.title, "title 4 - first duplicate")
}

优化探讨 #

首先明确：时间复杂度无法再低于O(n)，因为必须遍历数组中的每一个元素才能完成去重和分组逻辑。空间复杂度也已经是O(n)级别，因为需要存储结果（唯一元素和分组后的重复元素），这部分是无法避免的。

不过可以在细节上做优化，比如：

1. 减少内存拷贝：在实现一中，更新重复元素数组时，避免使用existingDuplicatesArray + [element]这种会创建新数组的方式，改用append直接修改原数组，减少不必要的内存分配。
2. 简化逻辑，利用Swift内置API：可以结合分组字典直接生成结果，避免二次遍历原数组：

public extension Array where Element: Identifiable {
    func filterDuplicates() -> (uniqueElements: [Element], duplicateElementGroups: [Element.ID: [Element]]) {
        let groupedById = Dictionary(grouping: self, by: \.id)
        // 唯一元素取每个分组的第一个元素
        let uniqueElements = groupedById.values.compactMap { $0.first }
        // 只保留元素数量>1的分组
        let duplicateGroups = groupedById.filter { $0.value.count > 1 }
        return (uniqueElements, duplicateGroups)
    }
}

这个版本利用Swift的高阶函数简化代码，同时保证逻辑清晰，时间复杂度同样是O(n)，空间上和实现二相当，但代码更简洁。需要注意的是，Swift 5.1+中Dictionary的遍历顺序是插入顺序，因此uniqueElements的顺序会和原数组中不同ID首次出现的顺序一致，和原有逻辑保持兼容。
3. 按需延迟计算：如果不需要同时持有唯一元素和重复元素集合，可以考虑将其中一个结果改为懒加载，但这取决于具体的使用场景。

总体来说，你的两种实现已经是非常高效的方案，上述优化更多是在代码简洁性和内存细节上的调整，不会改变整体的时间/空间复杂度级别。

内容的提问来源于stack exchange，提问作者Soheil Novinfard