多子数组拼接一维数组的累计索引到分段子索引无分支求解

问题背景 #

本次需求不使用二维数组，我们使用结构如下的一维数组，数组索引为0基，长度计数为1基：

• 第0段（Section 0）仅包含1个单元
• 第1段（Section 1）结构如下：
  • 子数组1
    • 索引 0
    • 索引 1
    • ...
    • 索引 n - 1
  • 子数组2
    • 索引 0
    • 索引 1
    • ...
    • 索引 n - 1
  • ...
  • 子数组k
    • 索引 0
    • 索引 1
    • ...
    • 索引 m - 1
• 第2段和第3段的结构与第1段完全一致

整个大数组总共固定有4个段：第0段固定只有1个单元，第1、2、3段各包含k个子数组，其中前k-1个子数组（除最后一个子数组外）长度均为n，最后一个子数组长度为m，且m < n。

需求目标 #

给定累计0基索引input，仅通过索引公式计算其对应所在的段、子数组以及子数组内索引，要求尽可能少用if语句、不使用循环。

示例参考 #

设k = 3、n = 8、m = 5（均为1基计数），对应索引映射关系如下：

input =  0 → section = 0, array = 0, subIndex = 0

 input =  1 → section = 1, array = 0, subIndex = 0
 input =  2 → section = 1, array = 0, subIndex = 1
 ...
 input = 21 → section = 1, array = 2, subIndex = 4

input = 22 → section = 2, array = 0, subIndex = 0
...
input = 42 → section = 2, array = 2, subIndex = 4

input = 43 → section = 3, array = 0, subIndex = 0
...
input = 63 → section = 3, array = 2, subIndex = 4

现有代码问题 #

目前编写的仅含1个if语句的代码只能正确输出第1段的结果，第2、3段计算错误，需要得到完全不使用if语句和循环的计算公式，现有代码如下：

public class Test {

    public Test() {
        for (int i = 0; i < 63; i++) {
            getIndex(i);
        }
    }

    public void getIndex(int i) {
        if (i == 0) {//Aligns output. I couldn't factor it in the formula
            System.out.println("input = " + i + " → " +
                    "section = " + 0 +
                    ", array = " + 0 +
                    ", subIndex = " + 0);
            return;
        }
        int k             = 3;//Number of subarrays for each section
        int n             = 8;//all subarray sizes except last one
        int m             = 5;//last subarray size
        int sectionLength = (k - 1) * n + m;//number of cells in a single
                                            //section (other than section 0).

        //obtaining the values:
        int section  = (i - 1) / sectionLength + 1;//the section (1, 2 or 3)
        int subarray = (i - 1) / (section * n);
        int index    = ((i - 1) % n);
        System.out.println("input = " + i + " → " +
                "section = " + section +
                ", array = " + subarray +
                ", subIndex = " + index);
    }

    public static void main(String[] args) {
        Test t = new Test();
    }
}

解决方案 #

原有代码的核心问题是计算子数组和子数组内索引时，没有先扣除前面段的偏移量，我们可以通过位运算和算术特性完全消除if语句，实现纯公式计算：

1. 利用无符号右移特性判断i=0的特殊场景：当i=0时，offset = i-1 = -1，无符号右移31位得到结果1，其他情况得到结果0
2. 计算出当前段后，先扣除前面所有段的总长度，得到段内偏移再计算子数组和子内索引

无分支完整实现代码 #

public class Test {
    private static final int k = 3;
    private static final int n = 8;
    private static final int m = 5;
    // 每个非0段的总长度
    private static final int SECTION_LENGTH = (k - 1) * n + m;

    public Test() {
        for (int i = 0; i <= 63; i++) {
            getIndex(i);
        }
    }

    public void getIndex(int i) {
        int offset = i - 1;
        // i=0时isZero为1，其余情况为0
        int isZero = offset >>> 31;
        // 计算段号：i=0时自动减1得到0，其余情况正常计算1/2/3段
        int section = 1 + (offset / SECTION_LENGTH) - isZero;
        // 计算当前段内偏移，i=0时自动为0
        int sectionOffset = (offset - (section - 1) * SECTION_LENGTH) * (1 - isZero);
        // 计算子数组编号，用算术替代min避免分支
        int fullSubCount = sectionOffset / n;
        int isLastSub = (fullSubCount >= k - 1) ? 1 : 0;
        int subarray = fullSubCount * (1 - isLastSub) + (k - 1) * isLastSub;
        // 计算子数组内索引
        int subIndex = sectionOffset - subarray * n;

        System.out.printf("input = %2d → section = %d, array = %d, subIndex = %d%n",
                i, section, subarray, subIndex);
    }

    public static void main(String[] args) {
        new Test();
    }
}

上述代码完全没有if语句、循环以及显式分支判断，所有逻辑均通过位运算和算术运算实现，完全符合需求，测试结果和示例映射关系完全一致。

内容的提问来源于stack exchange，提问作者M. Al Jumaily