erlang算法系列-线段树-2179. 统计数组中好三元组数目（困难）

2179. 统计数组中好三元组数目-原题

给你两个下标从 0 开始且长度为 n 的整数数组 nums1 和 nums2 ，两者都是 [0, 1, ..., n - 1] 的 排列 。

好三元组 指的是 3 个 互不相同 的值，且它们在数组 nums1 和 nums2 中出现顺序保持一致。换句话说，如果我们将 pos1v 记为值 v 在 nums1 中出现的位置，pos2v 为值 v 在 nums2 中的位置，那么一个好三元组定义为 0 <= x, y, z <= n - 1 ，且 pos1x < pos1y < pos1z 和 pos2x < pos2y < pos2z 都成立的 (x, y, z) 。

请你返回好三元组的 总数目 。

示例 1：

输入：nums1 = [2,0,1,3], nums2 = [0,1,2,3]
输出：1
解释：
总共有 4 个三元组 (x,y,z) 满足 pos1x < pos1y < pos1z ，分别是 (2,0,1) ，(2,0,3) ，(2,1,3) 和 (0,1,3) 。
这些三元组中，只有 (0,1,3) 满足 pos2x < pos2y < pos2z 。所以只有 1 个好三元组。
示例 2：

输入：nums1 = [4,0,1,3,2], nums2 = [4,1,0,2,3]
输出：4
解释：总共有 4 个好三元组 (4,0,3) ，(4,0,2) ，(4,1,3) 和 (4,1,2) 。

提示：

n == nums1.length == nums2.length
3 <= n <= 10^5
0 <= nums1[i], nums2[i] <= n - 1
nums1 和 nums2 是 [0, 1, ..., n - 1] 的排列。

解题思路

由于两个数组都是[0, 1, ..., n - 1] 的 排列，可以把题意稍微转化一下，先求出数组1中每个数分别在数组2的位置PosList。
例如nums1 = [2,0,1,3], nums2 = [0,1,2,3]，则PosList = [2,0,1,3]。
则结果为PosList中的每个数Pos,分别求出Pos左边小于Pos的个数Left，和Pos右边大于Pos的个数Right，所有Left*Right的和为解。
其中Right = Max - Pos -(I - Left), Max = n-1, I为Pos在PosList位置（从0开始）。

如何求Pos左边小于Pos的个数Left？使用线段树。

线段树结构：{Min, Max, Value, Small, Big} = TreeData。
Min数据段最小值，Max数据段最大值,Value数据段数据个数，Small左子树，Big右子树。
左右子树的划分值Mid = Min + (Max-Min) div 2, Small为{Min,Mid}，Big为{Mid+1,Max}。
每插入一个数Num的时候，更新所有Num属于{Min,Max}的节点的Value+1，时间复杂度O(lgN)。
同样计算所有{Min,Num-1}的个数，遍历线段树的时间复杂度也在O(lgN)。整体算法复杂度O(NlgN)。

-spec good_triplets(Nums1 :: [integer()], Nums2 :: [integer()]) -> integer().
good_triplets(Nums1, Nums2) ->
    Tree = gb_trees:empty(),
    NewTree = init_pos2(Nums2, 0, Tree),
    List = init_pos1(Nums1, NewTree, []),
    %Max - N -(I - Count - 1),
    Max = length(Nums1) - 1,
    TreeData = {0, Max, 0, nil, nil},    
    do_count(List, 0, Max, 1, TreeData, 0)
  .

init_pos1([], Tree, List) ->
    lists:reverse(List);
init_pos1([N | Nums], Tree, List) ->
    Pos = gb_trees:get(N, Tree),
    init_pos1(Nums, Tree, [Pos | List]).

init_pos2([], Ind, Tree) ->
    Tree;
init_pos2([N | Nums], Ind, Tree) ->
    NewTree = gb_trees:enter(N, Ind, Tree),
    init_pos2(Nums, Ind + 1, NewTree).

do_count([], Min, Max, Ind, Tree, Ans) ->
    Ans;
do_count([N | Nums], Min, Max, Ind, Tree, Ans) ->
    NewTree = insert_to_tree(Min, Max, N, Tree),
    Count = get_small_num(Min, Max, N, Tree),
    Count1 = Max - N - (Ind - Count - 1),
    do_count(Nums, Min, Max, Ind + 1, NewTree, Ans+ Count * Count1).


insert_to_tree(Min, Min, N, TreeData) ->
    case TreeData of
        {Min, Min, Value, nil, nil} ->
            {Min, Min, Value + 1, nil, nil};
        _ ->
            {Min, Min, 1, nil, nil}
    end;
insert_to_tree(Min, Max, N, TreeData) ->     
    if
        TreeData =:= nil ->
            Value = 0, Small = nil, Big = nil;
        true ->
            {Min, Max, Value, Small, Big} = TreeData
    end,    
    Mid = Min + (Max-Min) div 2,

    if
        Mid >= N ->
            {Min, Max, Value + 1, insert_to_tree(Min, Mid, N, Small), Big};
        true ->
            {Min, Max, Value + 1, Small, insert_to_tree(Mid + 1, Max, N, Big)}
    end. 

get_small_num(Min, Min, N, TreeData) ->
    case TreeData of
        {Min, Min, Value, nil, nil} ->
            if  
                N > Min ->
                    Value;
                true ->
                    0
            end;
        _ ->
            0
    end;
get_small_num(Min, Max, N, TreeData) ->     
    if
        TreeData =:= nil ->
            0;
        true ->
            {Min, Max, Value, Small, Big} = TreeData,
            if
                N > Max ->
                    Value;
                Min > N ->
                    0;
                true  ->
                    Mid = Min + (Max-Min) div 2,
                    if
                        Mid >= N ->
                            get_small_num(Min, Mid, N, Small);
                        true ->
                            get_small_num(Min, Mid, N, Small) + get_small_num(Mid + 1, Max, N, Big)
                    end
            end
    end.