667. Beautiful Arrangement II

Given two integers n and k, you need to construct a list which contains n different positive integers ranging from 1 to n and obeys the following requirement:
Suppose this list is [a1, a2, a3, ... , an], then the list [|a1 - a2|, |a2 - a3|, |a3 - a4|, ... , |an-1 - an|] has exactly k distinct integers.

If there are multiple answers, print any of them.

Example 1:

Input: n = 3, k = 1
Output: [1, 2, 3]
Explanation: The [1, 2, 3] has three different positive integers ranging from 1 to 3, and the [1, 1] has exactly 1 distinct integer: 1.

Example 2:

Input: n = 3, k = 2
Output: [1, 3, 2]
Explanation: The [1, 3, 2] has three different positive integers ranging from 1 to 3, and the [2, 1] has exactly 2 distinct integers: 1 and 2.

Note:

  1. The n and k are in the range 1 <= k < n <= 104.

Rust Solution

struct Solution;

impl Solution {
    fn construct_array(n: i32, mut k: i32) -> Vec<i32> {
        let mut res = vec![1];
        let mut l = 2;
        let mut r = n;
        let mut forward = true;
        for _ in 1..n {
            if k > 1 {
                if forward {
                    res.push(r);
                    r -= 1;
                } else {
                    res.push(l);
                    l += 1;
                }
                forward = !forward;
                k -= 1;
            } else {
                if forward {
                    res.push(l);
                    l += 1;
                } else {
                    res.push(r);
                    r -= 1;
                }
            }
        }
        res
    }
}

#[test]
fn test() {
    let n = 3;
    let k = 1;
    let res = vec![1, 2, 3];
    assert_eq!(Solution::construct_array(n, k), res);
    let n = 3;
    let k = 2;
    let res = vec![1, 3, 2];
    assert_eq!(Solution::construct_array(n, k), res);
    let n = 5;
    let k = 2;
    let res = vec![1, 5, 4, 3, 2];
    assert_eq!(Solution::construct_array(n, k), res);
}

Having problems with this solution? Click here to submit an issue on github.