057_top_k_frequent.sio

//@ run-pass
// HumanEval 057: Top K Frequent Elements
//
// Find the k most frequent elements in an array. Uses a frequency
// count approach: first count occurrences of each unique value, then
// pick the k values with the highest counts via partial selection sort.
// Returns unique values sorted by frequency (descending).

struct Buf { data: [i64; 256] }

fn top_k_frequent(arr: [i64; 256], n: i64, k: i64, out: &!Buf) -> i64 with Mut, Panic {
    // Step 1: Count frequencies of unique values
    var vals: [i64; 256] = [0; 256]
    var counts: [i64; 256] = [0; 256]
    var num_unique: i64 = 0

    var i: i64 = 0
    while i < n {
        // Search for arr[i] in vals
        var found: i64 = 0
        var j: i64 = 0
        while j < num_unique {
            if vals[j] == arr[i] {
                counts[j] = counts[j] + 1
                found = 1
                break
            }
            j = j + 1
        }
        if found == 0 {
            vals[num_unique] = arr[i]
            counts[num_unique] = 1
            num_unique = num_unique + 1
        }
        i = i + 1
    }

    // Step 2: Partial selection sort by count (descending), pick top k
    var pass: i64 = 0
    while pass < k {
        if pass >= num_unique { break }
        var max_idx = pass
        var j: i64 = pass + 1
        while j < num_unique {
            if counts[j] > counts[max_idx] {
                max_idx = j
            }
            j = j + 1
        }
        // Swap vals and counts
        let tv = vals[pass]
        vals[pass] = vals[max_idx]
        vals[max_idx] = tv
        let tc = counts[pass]
        counts[pass] = counts[max_idx]
        counts[max_idx] = tc

        out.data[pass] = vals[pass]
        pass = pass + 1
    }

    if k < num_unique { k } else { num_unique }
}

fn main() -> i64 with IO, Mut, Panic, Div {
    // Test 1: [1,1,1,2,2,3], k=2 -> [1, 2] (1 appears 3x, 2 appears 2x)
    var a1: [i64; 256] = [0; 256]
    a1[0] = 1
    a1[1] = 1
    a1[2] = 1
    a1[3] = 2
    a1[4] = 2
    a1[5] = 3
    var o1 = Buf { data: [0; 256] }
    let n1 = top_k_frequent(a1, 6, 2, &!o1)
    assert(n1 == 2)
    assert(o1.data[0] == 1)
    assert(o1.data[1] == 2)

    // Test 2: [1], k=1 -> [1]
    var a2: [i64; 256] = [0; 256]
    a2[0] = 1
    var o2 = Buf { data: [0; 256] }
    let n2 = top_k_frequent(a2, 1, 1, &!o2)
    assert(n2 == 1)
    assert(o2.data[0] == 1)

    // Test 3: [4,4,4,5,5,6,6,6,6], k=1 -> [6] (6 appears 4x)
    var a3: [i64; 256] = [0; 256]
    a3[0] = 4
    a3[1] = 4
    a3[2] = 4
    a3[3] = 5
    a3[4] = 5
    a3[5] = 6
    a3[6] = 6
    a3[7] = 6
    a3[8] = 6
    var o3 = Buf { data: [0; 256] }
    let n3 = top_k_frequent(a3, 9, 1, &!o3)
    assert(n3 == 1)
    assert(o3.data[0] == 6)

    // Test 4: [10,20,10,30,20,10], k=3 -> [10, 20, 30]
    var a4: [i64; 256] = [0; 256]
    a4[0] = 10
    a4[1] = 20
    a4[2] = 10
    a4[3] = 30
    a4[4] = 20
    a4[5] = 10
    var o4 = Buf { data: [0; 256] }
    let n4 = top_k_frequent(a4, 6, 3, &!o4)
    assert(n4 == 3)
    assert(o4.data[0] == 10)
    assert(o4.data[1] == 20)
    assert(o4.data[2] == 30)

    println("057_top_k_frequent: ALL TESTS PASSED")
    0
}