如何编写高效awk函数实现带来源标注的有序数组合并

问题需求 #

给定两个升序排列的数组，需要将它们合并为一个有序数组，同时用另一个数组记录每个元素的来源（1代表来自第一个输入数组，2代表来自第二个）。

示例输入 #

x1[1]=10
x1[2]=20
x1[3]=30

x2[1]=15
x2[2]=25
x2[3]=35

期望输出 #

x[1]=10
x[2]=15
x[3]=20
x[4]=25
x[5]=30
x[6]=35
a[1]=1
a[2]=2
a[3]=1
a[4]=2
a[5]=1
a[6]=2

需要编写一个高效的AWK函数完成这个任务，函数框架如下：

function f(r, a, x1, x2) {
# r output merge array
# a an annotation array indicating whether an element at a given index is from x1 or x2
# x1,x2 input sorted arrays
  ...
}

高效AWK实现函数 #

这里采用双指针法实现，时间复杂度为O(n+m)（n和m分别是两个输入数组的长度），是合并有序数组的最优效率方案：

function f(r, a, x1, x2,    i, j, k) {
    i = j = k = 1
    # 同时遍历两个数组，取较小的元素加入结果
    while (i in x1 && j in x2) {
        if (x1[i] <= x2[j]) {
            r[k] = x1[i]
            a[k++] = 1
            i++
        } else {
            r[k] = x2[j]
            a[k++] = 2
            j++
        }
    }
    # 处理第一个数组剩下的元素
    while (i in x1) {
        r[k] = x1[i]
        a[k++] = 1
        i++
    }
    # 处理第二个数组剩下的元素
    while (j in x2) {
        r[k] = x2[j]
        a[k++] = 2
        j++
    }
}

函数逻辑说明 #

• 用i和j分别作为两个输入数组的遍历指针，k作为结果数组的索引指针，初始都设为1（符合AWK数组通常从1开始的习惯）。
• 第一循环：比较当前两个指针指向的元素，将较小的那个加入结果数组r，同时在a数组标记来源，然后移动对应指针。
• 当其中一个数组遍历完后，将另一个数组剩余的元素直接追加到结果数组末尾，同时标记来源。
• 函数内部的i,j,k是局部变量（通过函数参数列表末尾的逗号后声明，AWK中这类变量为函数内局部变量），避免污染全局变量。

测试示例 #

可以用以下代码测试这个函数：

BEGIN {
    # 初始化输入数组
    x1[1]=10; x1[2]=20; x1[3]=30
    x2[1]=15; x2[2]=25; x2[3]=35
    
    # 调用合并函数
    f(x, a, x1, x2)
    
    # 输出结果数组x
    for (k=1; k<=length(x); k++) {
        printf "x[%d]=%d\n", k, x[k]
    }
    # 输出来源数组a
    for (k=1; k<=length(a); k++) {
        printf "a[%d]=%d\n", k, a[k]
    }
}

# 插入上面的f函数
function f(r, a, x1, x2,    i, j, k) {
    i = j = k = 1
    while (i in x1 && j in x2) {
        if (x1[i] <= x2[j]) {
            r[k] = x1[i]
            a[k++] = 1
            i++
        } else {
            r[k] = x2[j]
            a[k++] = 2
            j++
        }
    }
    while (i in x1) {
        r[k] = x1[i]
        a[k++] = 1
        i++
    }
    while (j in x2) {
        r[k] = x2[j]
        a[k++] = 2
        j++
    }
}

运行这段AWK代码，就能得到期望的输出结果。

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